// graph-tool -- a general graph modification and manipulation thingy
//
// Copyright (C) 2006-2018 Tiago de Paula Peixoto <tiago@skewed.de>
//
// This program is free software; you can redistribute it and/or
// modify it under the terms of the GNU General Public License
// as published by the Free Software Foundation; either version 3
// of the License, or (at your option) any later version.
//
// 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
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.

#ifndef GRAPH_REWIRING_HH
#define GRAPH_REWIRING_HH

#include <tuple>
#include <iostream>
#include <boost/functional/hash.hpp>

#include "graph.hh"
#include "graph_filtering.hh"
#include "graph_util.hh"
#include "sampler.hh"

#include "random.hh"

#include "hash_map_wrap.hh"

namespace graph_tool
{
using namespace std;
using namespace boost;

template <class Graph>
typename graph_traits<Graph>::vertex_descriptor
source(const pair<size_t, bool>& e,
       const vector<typename graph_traits<Graph>::edge_descriptor>& edges,
       const Graph& g)
{
    if (e.second)
        return target(edges[e.first], g);
    else
        return source(edges[e.first], g);
}

template <class Graph>
typename graph_traits<Graph>::vertex_descriptor
target(const pair<size_t, bool>& e,
       const vector<typename graph_traits<Graph>::edge_descriptor>& edges,
       const Graph& g)
{
    if (e.second)
        return source(edges[e.first], g);
    else
        return target(edges[e.first], g);
}

template <class Nmap, class Graph>
void add_count(size_t s, size_t t, Nmap& nvmap, Graph& g)
{
    if (!graph_tool::is_directed(g) && s > t)
        std::swap(s, t);
    auto& nmap = nvmap[s];
    nmap[t]++;
}

template <class Nmap, class Graph>
void remove_count(size_t s, size_t t, Nmap& nvmap, Graph& g)
{
    if (!graph_tool::is_directed(g) && s > t)
        std::swap(s, t);
    auto& nmap = nvmap[s];
    auto iter = nmap.find(t);
    iter->second--;
    if (iter->second == 0)
        nmap.erase(iter);
}

template <class Nmap, class Graph>
size_t get_count(size_t s, size_t t, Nmap& nvmap, Graph& g)
{
    if (!graph_tool::is_directed(g) && s > t)
        std::swap(s, t);
    auto& nmap = nvmap[s];
    auto iter = nmap.find(t);
    if (iter == nmap.end())
        return 0;
    return iter->second;
    // if (s != t)
    //     return iter->second;
    // else
    //     return iter->second / 2;
}

// this functor will swap the source of the edge e with the source of edge se
// and the target of edge e with the target of te
struct swap_edge
{
    template <class Nmap, class Graph>
    static bool
    parallel_check_target (const pair<size_t, bool>& e,
                           const pair<size_t, bool>& te,
                           vector<typename graph_traits<Graph>::edge_descriptor>& edges,
                           Nmap& nmap,
                           const Graph &g)
    {
        // We want to check that if we swap the target of 'e' with the target of
        // 'te', as such
        //
        //  (s)    -e--> (t)          (s)    -e--> (nt)
        //  (te_s) -te-> (nt)   =>    (te_s) -te-> (t)
        //
        // no parallel edges are introduced.

        typename graph_traits<Graph>::vertex_descriptor
            s = source(e, edges, g),          // current source
            t = target(e, edges, g),          // current target
            nt = target(te, edges, g),        // new target
            te_s = source(te, edges, g);      // target edge source

        if (get_count(s,  nt, nmap, g) > 0)
            return true; // e would clash with an existing edge
        if (get_count(te_s, t, nmap, g) > 0)
            return true; // te would clash with an existing edge
        return false; // no parallel edges
    }

    template <class Graph>
    static void swap_target
        (const pair<size_t, bool>& e,
         const pair<size_t, bool>& te,
         vector<typename graph_traits<Graph>::edge_descriptor>& edges,
         Graph& g)
    {
        // swap the target of the edge 'e' with the target of edge 'te', as
        // such:
        //
        //  (s)    -e--> (t)          (s)    -e--> (nt)
        //  (te_s) -te-> (nt)   =>    (te_s) -te-> (t)

        if (e.first == te.first)
            return;

        // new edges which will replace the old ones
        typename graph_traits<Graph>::edge_descriptor ne, nte;
        typename graph_traits<Graph>::vertex_descriptor
            s_e = source(e, edges, g),
            t_e = target(e, edges, g),
            s_te = source(te, edges, g),
            t_te = target(te, edges, g);
        remove_edge(edges[e.first], g);
        remove_edge(edges[te.first], g);

        if (graph_tool::is_directed(g) || !e.second)
            ne = add_edge(s_e, t_te, g).first;
        else // keep invertedness (only for undirected graphs)
            ne = add_edge(t_te, s_e, g).first;
        edges[e.first] = ne;
        if (graph_tool::is_directed(g) || !te.second)
            nte = add_edge(s_te, t_e, g).first;
        else // keep invertedness (only for undirected graphs)
            nte = add_edge(t_e, s_te,  g).first;
        edges[te.first] = nte;
    }

};

// used for verbose display
void print_progress(size_t i, size_t n_iter, size_t current, size_t total,
                    stringstream& str)
{
    size_t atom = (total > 200) ? total / 100 : 1;
    if ( ( (current+1) % atom == 0) || (current + 1) == total)
    {
        size_t size = str.str().length();
        for (size_t j = 0; j < str.str().length(); ++j)
            cout << "\b";
        str.str("");
        str << "(" << i + 1 << " / " << n_iter << ") "
            << current + 1 << " of " << total << " ("
            << (current + 1) * 100 / total << "%)";
        for (int j = 0; j < int(size - str.str().length()); ++j)
            str << " ";
        cout << str.str() << flush;
    }
}

//select blocks based on in/out degrees
class DegreeBlock
{
public:
    typedef pair<size_t, size_t> block_t;

    template <class Graph>
    block_t get_block(typename graph_traits<Graph>::vertex_descriptor v,
                      const Graph& g) const
    {
        return make_pair(in_degreeS()(v, g), out_degree(v, g));
    }
};

//select blocks based on property map
template <class PropertyMap>
class PropertyBlock
{
public:
    typedef typename property_traits<PropertyMap>::value_type block_t;

    PropertyBlock(PropertyMap p): _p(p) {}

    template <class Graph>
    block_t get_block(typename graph_traits<Graph>::vertex_descriptor v,
                      const Graph&) const
    {
        return get(_p, v);
    }

private:
    PropertyMap _p;
};

// select an appropriate "null" key for densehash
template <class Type>
struct get_null_key
{
    Type operator()() const
    {
        return numeric_limits<Type>::max();
    }
};

template <>
struct get_null_key<string>
{
    string operator()() const
    {
        return lexical_cast<string>(get_null_key<size_t>()());
    }
};

template <>
struct get_null_key<boost::python::object>
{
    boost::python::object operator()() const
    {
        return boost::python::object();
    }
};

template <class Type>
struct get_null_key<vector<Type>>
{
    vector<Type> operator()() const
    {
        vector<Type> v(1);
        v[0] = get_null_key<Type>()();
        return v;
    }
};

template <class Type1, class Type2>
struct get_null_key<pair<Type1, Type2>>
{
    pair<Type1, Type2> operator()() const
    {
        return make_pair(get_null_key<Type1>()(),
                         get_null_key<Type2>()());
    }
};

// main rewire loop
template <template <class Graph, class EdgeIndexMap, class CorrProb,
                    class BlockDeg>
          class RewireStrategy>
struct graph_rewire
{

    template <class Graph, class EdgeIndexMap, class CorrProb,
              class BlockDeg, class PinMap>
    void operator()(Graph& g, EdgeIndexMap edge_index, CorrProb corr_prob,
                    PinMap pin, bool self_loops, bool parallel_edges,
                    bool configuration, pair<size_t, bool> iter_sweep,
                    std::tuple<bool, bool, bool> cache_verbose, size_t& pcount,
                    rng_t& rng, BlockDeg bd) const
    {
        typedef typename graph_traits<Graph>::edge_descriptor edge_t;
        bool persist = std::get<0>(cache_verbose);
        bool cache = std::get<1>(cache_verbose);
        bool verbose = std::get<2>(cache_verbose);

        vector<edge_t> edges;
        vector<size_t> edge_pos;
        typename graph_traits<Graph>::edge_iterator e, e_end;
        for (tie(e, e_end) = boost::edges(g); e != e_end; ++e)
        {
            if (pin[*e])
                continue;
            edges.push_back(*e);
            edge_pos.push_back(edge_pos.size());
        }

        typedef random_permutation_iterator<typename vector<size_t>::iterator,
                                            rng_t>
            random_edge_iter;

        RewireStrategy<Graph, EdgeIndexMap, CorrProb, BlockDeg>
            rewire(g, edge_index, edges, corr_prob, bd, cache, rng,
                   parallel_edges, configuration);

        size_t niter;
        bool no_sweep;
        tie(niter, no_sweep) = iter_sweep;
        pcount = 0;
        if (verbose)
            cout << "rewiring edges: ";
        stringstream str;
        for (size_t i = 0; i < niter; ++i)
        {
            random_edge_iter
                ei_begin(edge_pos.begin(), edge_pos.end(), rng),
                ei_end(edge_pos.end(), edge_pos.end(), rng);

            for (random_edge_iter ei = ei_begin; ei != ei_end; ++ei)
            {
                size_t e_pos = ei - ei_begin;
                if (verbose)
                    print_progress(i, niter, e_pos, no_sweep ? 1 : edges.size(),
                                   str);

                size_t e = *ei;

                bool success = false;
                do
                {
                    success = rewire(e, self_loops, parallel_edges);
                }
                while(persist && !success);

                if (!success)
                    ++pcount;

                if (no_sweep)
                    break;
            }
        }
        if (verbose)
            cout << endl;
    }

    template <class Graph, class EdgeIndexMap, class CorrProb, class PinMap>
    void operator()(Graph& g, EdgeIndexMap edge_index, CorrProb corr_prob,
                    PinMap pin, bool self_loops, bool parallel_edges,
                    bool configuration, pair<size_t, bool> iter_sweep,
                    std::tuple<bool, bool, bool> cache_verbose, size_t& pcount,
                    rng_t& rng) const
    {
        operator()(g, edge_index, corr_prob, pin, self_loops, parallel_edges,
                   configuration, iter_sweep, cache_verbose, pcount, rng,
                   DegreeBlock());
    }
};

// this will rewire the edges so that the resulting graph will be entirely
// random (i.e. Erdos-Renyi)
template <class Graph, class EdgeIndexMap, class CorrProb, class BlockDeg>
class ErdosRewireStrategy
{
public:
    typedef typename graph_traits<Graph>::edge_descriptor edge_t;
    typedef typename EdgeIndexMap::value_type index_t;

    ErdosRewireStrategy(Graph& g, EdgeIndexMap edge_index,
                        vector<edge_t>& edges, CorrProb, BlockDeg,
                        bool, rng_t& rng, bool, bool configuration)
        : _g(g), _edge_index(edge_index), _edges(edges),
          _vertices(HardNumVertices()(g)), _rng(rng),
          _configuration(configuration),
          _nmap(get(vertex_index, g), num_vertices(g))
    {
        decltype(_vertices.begin()) viter = _vertices.begin();
        typename graph_traits<Graph>::vertex_iterator v, v_end;
        for (tie(v, v_end) = vertices(_g); v != v_end; ++v)
            *(viter++) = *v;

        if (!configuration)
        {
            for (size_t i = 0; i < edges.size(); ++i)
                add_count(source(edges[i], g), target(edges[i], g), _nmap, g);
        }
    }

    bool operator()(size_t ei, bool self_loops, bool parallel_edges)
    {
        size_t e_s = source(_edges[ei], _g);
        size_t e_t = target(_edges[ei], _g);

        if (!graph_tool::is_directed(_g) && e_s > e_t)
            std::swap(e_s, e_t);

        //try randomly drawn pairs of vertices
        std::uniform_int_distribution<size_t> sample(0, _vertices.size() - 1);
        typename graph_traits<Graph>::vertex_descriptor s, t;
        while (true)
        {
            s = sample(_rng);
            t = sample(_rng);
            if (s == t)
            {
                if (!self_loops) // reject self-loops if not allowed
                    continue;

            }
            else if (!graph_tool::is_directed(_g) && self_loops)
            {
                // sample self-loops w/ correct probability for undirected
                // graphs
                std::bernoulli_distribution reject(.5);
                if (reject(_rng))
                    continue;
            }
            break;
        }

        if (!graph_tool::is_directed(_g) && s > t)
            std::swap(s, t);

        if (s == e_s && t == e_t)
            return false;

        // reject parallel edges if not allowed
        if (!parallel_edges && is_adjacent(s, t, _g))
            return false;

        if (!_configuration)
        {
            size_t m = get_count(s, t, _nmap, _g);
            size_t m_e = get_count(e_s, e_t, _nmap, _g);

            double a = (m + 1) / double(m_e);

            std::bernoulli_distribution accept(std::min(a, 1.));
            if (!accept(_rng))
                return false;
        }

        remove_edge(_edges[ei], _g);
        edge_t ne = add_edge(s, t, _g).first;
        _edges[ei] = ne;

        if (!_configuration)
        {
            remove_count(e_s, e_t, _nmap, _g);
            add_count(s, t, _nmap, _g);
        }

        return true;
    }

private:
    Graph& _g;
    EdgeIndexMap _edge_index;
    vector<edge_t>& _edges;
    vector<typename graph_traits<Graph>::vertex_descriptor> _vertices;
    rng_t& _rng;
    bool _configuration;
    typedef gt_hash_map<size_t, size_t> nmapv_t;
    typedef typename property_map_type::apply<nmapv_t,
                                              typename property_map<Graph, vertex_index_t>::type>
        ::type::unchecked_t nmap_t;
    nmap_t _nmap;
};

// this is the mother class for edge-based rewire strategies
// it contains the common loop for finding edges to swap, so different
// strategies need only to specify where to sample the edges from.
template <class Graph, class EdgeIndexMap, class RewireStrategy>
class RewireStrategyBase
{
public:
    typedef typename graph_traits<Graph>::edge_descriptor edge_t;
    typedef typename graph_traits<Graph>::vertex_descriptor vertex_t;

    typedef typename EdgeIndexMap::value_type index_t;

    RewireStrategyBase(Graph& g, EdgeIndexMap edge_index, vector<edge_t>& edges,
                       rng_t& rng, bool parallel_edges, bool configuration)
        : _g(g), _edge_index(edge_index), _edges(edges), _rng(rng),
          _nmap(get(vertex_index, g), num_vertices(g)),
          _configuration(configuration)
    {
        if (!parallel_edges || !configuration)
        {
            for (size_t i = 0; i < edges.size(); ++i)
                add_count(source(edges[i], g), target(edges[i], g), _nmap, g);
        }
    }

    bool operator()(size_t ei, bool self_loops, bool parallel_edges)
    {
        RewireStrategy& self = *static_cast<RewireStrategy*>(this);

        // try randomly drawn pairs of edges and check if they satisfy all the
        // consistency checks

        pair<size_t, bool> e = make_pair(ei, false);

        // rewire target
        pair<size_t, bool> et = self.get_target_edge(e, parallel_edges);

        if (et.first == ei)
            return false;

        auto s = source(e,  _edges, _g);
        auto t = target(e,  _edges, _g);
        auto ts = source(et,  _edges, _g);
        auto tt = target(et,  _edges, _g);

        if (!self_loops) // reject self-loops if not allowed
        {
            if(s == tt || ts == t)
                return false;
        }

        // reject parallel edges if not allowed
        if (!parallel_edges && (et.first != e.first))
        {
            if (swap_edge::parallel_check_target(e, et, _edges, _nmap, _g))
                return false;
        }

        double a = 0;

        if (!graph_tool::is_directed(_g))
        {
            a -= log(2 + (s == t) + (ts == tt));
            a += log(2 + (s == tt) + (ts == t));
        }

        if (!_configuration)
        {
            map<std::pair<size_t, size_t>, int> delta;

            delta[std::make_pair(s, t)] -= 1;
            delta[std::make_pair(ts, tt)] -= 1;
            delta[std::make_pair(s, tt)] += 1;
            delta[std::make_pair(ts, t)] += 1;

            for (auto& e_d : delta)
            {
                auto u = e_d.first.first;
                auto v = e_d.first.second;
                int d = e_d.second;
                size_t m = get_count(u,  v,  _nmap, _g);
                a -= lgamma(m + 1) - lgamma((m + 1) + d);
                if (!graph_tool::is_directed(_g) && u == v)
                    a += d * log(2);
            }

        }

        std::bernoulli_distribution accept(std::min(exp(a), 1.));
        if (!accept(_rng))
            return false;

        self.update_edge(e.first, false);
        self.update_edge(et.first, false);

        if (!parallel_edges || !_configuration)
        {
            remove_count(source(e, _edges, _g), target(e, _edges, _g), _nmap, _g);
            remove_count(source(et, _edges, _g), target(et, _edges, _g), _nmap, _g);
        }

        swap_edge::swap_target(e, et, _edges, _g);

        self.update_edge(e.first, true);
        self.update_edge(et.first, true);

        if (!parallel_edges || !_configuration)
        {
            add_count(source(e, _edges, _g), target(e, _edges, _g), _nmap, _g);
            add_count(source(et, _edges, _g), target(et, _edges, _g), _nmap, _g);
        }

        return true;
    }

protected:
    Graph& _g;
    EdgeIndexMap _edge_index;
    vector<edge_t>& _edges;
    rng_t& _rng;

    typedef gt_hash_map<size_t, size_t> nmapv_t;
    typedef typename property_map_type::apply<nmapv_t,
                                              typename property_map<Graph, vertex_index_t>::type>
        ::type::unchecked_t nmap_t;
    nmap_t _nmap;
    bool _configuration;
};

// this will rewire the edges so that the combined (in, out) degree distribution
// will be the same, but all the rest is random
template <class Graph, class EdgeIndexMap, class CorrProb, class BlockDeg>
class RandomRewireStrategy:
    public RewireStrategyBase<Graph, EdgeIndexMap,
                              RandomRewireStrategy<Graph, EdgeIndexMap,
                                                   CorrProb, BlockDeg> >
{
public:
    typedef RewireStrategyBase<Graph, EdgeIndexMap,
                               RandomRewireStrategy<Graph, EdgeIndexMap,
                                                    CorrProb, BlockDeg> >
        base_t;

    typedef Graph graph_t;
    typedef EdgeIndexMap edge_index_t;

    typedef typename graph_traits<Graph>::vertex_descriptor vertex_t;
    typedef typename graph_traits<Graph>::edge_descriptor edge_t;
    typedef typename EdgeIndexMap::value_type index_t;

    struct hash_index {};
    struct random_index {};

    RandomRewireStrategy(Graph& g, EdgeIndexMap edge_index,
                         vector<edge_t>& edges, CorrProb, BlockDeg,
                         bool, rng_t& rng, bool parallel_edges,
                         bool configuration)
        : base_t(g, edge_index, edges, rng, parallel_edges, configuration),
          _g(g) {}

    pair<size_t,bool> get_target_edge(pair<size_t,bool>& e, bool)
    {
        std::uniform_int_distribution<> sample(0, base_t::_edges.size() - 1);
        pair<size_t, bool> et = make_pair(sample(base_t::_rng), false);
        if (!graph_tool::is_directed(_g))
        {
            std::bernoulli_distribution coin(0.5);
            et.second = coin(base_t::_rng);
            e.second = coin(base_t::_rng);
        }
        return et;
    }

    void update_edge(size_t, bool) {}

private:
    Graph& _g;
    EdgeIndexMap _edge_index;
};

// this will rewire the edges so that the (in,out) degree distributions and the
// (in,out)->(in,out) correlations will be the same, but all the rest is random
template <class Graph, class EdgeIndexMap, class CorrProb, class BlockDeg>
class CorrelatedRewireStrategy:
    public RewireStrategyBase<Graph, EdgeIndexMap,
                              CorrelatedRewireStrategy<Graph, EdgeIndexMap,
                                                       CorrProb, BlockDeg> >
{
public:
    typedef RewireStrategyBase<Graph, EdgeIndexMap,
                               CorrelatedRewireStrategy<Graph, EdgeIndexMap,
                                                        CorrProb, BlockDeg> >
        base_t;

    typedef Graph graph_t;

    typedef typename graph_traits<Graph>::vertex_descriptor vertex_t;
    typedef typename graph_traits<Graph>::edge_descriptor edge_t;

    typedef typename BlockDeg::block_t deg_t;

    CorrelatedRewireStrategy(Graph& g, EdgeIndexMap edge_index,
                             vector<edge_t>& edges, CorrProb, BlockDeg blockdeg,
                             bool, rng_t& rng, bool parallel_edges,
                             bool configuration)
        : base_t(g, edge_index, edges, rng, parallel_edges, configuration),
          _blockdeg(blockdeg), _g(g)
    {
        for (size_t ei = 0; ei < base_t::_edges.size(); ++ei)
        {
            // For undirected graphs, there is no difference between source and
            // target, and each edge will appear _twice_ in the list below,
            // once for each different ordering of source and target.
            edge_t& e = base_t::_edges[ei];

            vertex_t t = target(e, _g);
            deg_t tdeg = get_deg(t, _g);;
            _edges_by_target[tdeg].push_back(make_pair(ei, false));

            if (!graph_tool::is_directed(_g))
            {
                t = source(e, _g);
                deg_t tdeg = get_deg(t, _g);
                _edges_by_target[tdeg].push_back(make_pair(ei, true));
            }
        }
    }

    pair<size_t,bool> get_target_edge(pair<size_t, bool>& e, bool)
    {
        if (!graph_tool::is_directed(_g))
        {
            std::bernoulli_distribution coin(0.5);
            e.second = coin(base_t::_rng);
        }

        vertex_t t = target(e, base_t::_edges, _g);
        deg_t tdeg = get_deg(t, _g);
        auto& elist = _edges_by_target[tdeg];
        std::uniform_int_distribution<> sample(0, elist.size() - 1);
        auto ep = elist[sample(base_t::_rng)];
        if (get_deg(target(ep, base_t::_edges, _g), _g) != tdeg)
            ep.second = not ep.second;
        return ep;
    }

    void update_edge(size_t, bool) {}

    deg_t get_deg(vertex_t v, const Graph& g)
    {
        return _blockdeg.get_block(v, g);
    }

private:
    BlockDeg _blockdeg;

    typedef std::unordered_map<deg_t,
                          vector<pair<size_t, bool>>>
        edges_by_end_deg_t;

    edges_by_end_deg_t _edges_by_target;

protected:
    const Graph& _g;
};

// general stochastic blockmodel
// this version is based on rejection sampling
template <class Graph, class EdgeIndexMap, class CorrProb, class BlockDeg>
class ProbabilisticRewireStrategy:
    public RewireStrategyBase<Graph, EdgeIndexMap,
                              ProbabilisticRewireStrategy<Graph, EdgeIndexMap,
                                                          CorrProb, BlockDeg> >
{
public:
    typedef RewireStrategyBase<Graph, EdgeIndexMap,
                               ProbabilisticRewireStrategy<Graph, EdgeIndexMap,
                                                           CorrProb, BlockDeg> >
        base_t;

    typedef Graph graph_t;
    typedef EdgeIndexMap edge_index_t;

    typedef typename BlockDeg::block_t deg_t;

    typedef typename graph_traits<Graph>::vertex_descriptor vertex_t;
    typedef typename graph_traits<Graph>::edge_descriptor edge_t;
    typedef typename EdgeIndexMap::value_type index_t;

    ProbabilisticRewireStrategy(Graph& g, EdgeIndexMap edge_index,
                                vector<edge_t>& edges, CorrProb corr_prob,
                                BlockDeg blockdeg, bool cache, rng_t& rng,
                                bool parallel_edges, bool configuration)
        : base_t(g, edge_index, edges, rng, parallel_edges, configuration),
          _g(g), _corr_prob(corr_prob), _blockdeg(blockdeg)
    {
        if (cache)
        {
            // cache probabilities
            _corr_prob.get_probs(_probs);

            if (_probs.empty())
            {
                std::unordered_set<deg_t> deg_set;
                for (size_t ei = 0; ei < base_t::_edges.size(); ++ei)
                {
                    edge_t& e = base_t::_edges[ei];
                    deg_set.insert(get_deg(source(e, g), g));
                    deg_set.insert(get_deg(target(e, g), g));
                }

                for (auto s_iter = deg_set.begin(); s_iter != deg_set.end(); ++s_iter)
                    for (auto t_iter = deg_set.begin(); t_iter != deg_set.end(); ++t_iter)
                    {
                        double p = _corr_prob(*s_iter, *t_iter);
                        _probs[make_pair(*s_iter, *t_iter)] = p;
                    }
            }

            for (auto iter = _probs.begin(); iter != _probs.end(); ++iter)
            {
                double& p = iter->second;
                // avoid zero probability to not get stuck in rejection step
                if (std::isnan(p) || std::isinf(p) || p <= 0)
                    p = numeric_limits<double>::min();
                p = log(p);
            }
        }
    }

    double get_prob(const deg_t& s_deg, const deg_t& t_deg)
    {
        if (_probs.empty())
        {
            double p = _corr_prob(s_deg, t_deg);
            // avoid zero probability to not get stuck in rejection step
            if (std::isnan(p) || std::isinf(p) || p <= 0)
                p = numeric_limits<double>::min();
            return log(p);
        }
        auto k = make_pair(s_deg, t_deg);
        auto iter = _probs.find(k);
        if (iter == _probs.end())
            return log(numeric_limits<double>::min());
        return iter->second;
    }

    deg_t get_deg(vertex_t v, Graph& g)
    {
        return _blockdeg.get_block(v, g);
    }

    pair<size_t, bool> get_target_edge(pair<size_t, bool>& e, bool)
    {
        if (!graph_tool::is_directed(_g))
        {
            std::bernoulli_distribution coin(0.5);
            e.second = coin(base_t::_rng);
        }

        deg_t s_deg = get_deg(source(e, base_t::_edges, _g), _g);
        deg_t t_deg = get_deg(target(e, base_t::_edges, _g), _g);

        std::uniform_int_distribution<> sample(0, base_t::_edges.size() - 1);
        size_t epi = sample(base_t::_rng);
        pair<size_t, bool> ep = make_pair(epi, false);
        if (!graph_tool::is_directed(_g))
        {
            // for undirected graphs we must select a random direction
            std::bernoulli_distribution coin(0.5);
            ep.second = coin(base_t::_rng);
        }

        if (source(e, base_t::_edges, _g) == source(ep, base_t::_edges, _g) ||
            target(e, base_t::_edges, _g) == target(ep, base_t::_edges, _g))
            return ep; // rewiring is a no-op

        deg_t ep_s_deg = get_deg(source(ep, base_t::_edges, _g), _g);
        deg_t ep_t_deg = get_deg(target(ep, base_t::_edges, _g), _g);

        double pi = get_prob(s_deg, t_deg) + get_prob(ep_s_deg, ep_t_deg);
        double pf = get_prob(s_deg, ep_t_deg) + get_prob(ep_s_deg, t_deg);

        if (pf >= pi)
            return ep;

        double a = exp(pf - pi);

        std::uniform_real_distribution<> rsample(0.0, 1.0);
        double r = rsample(base_t::_rng);
        if (r > a)
            return e; // reject
        else
            return ep;
    }

    void update_edge(size_t, bool) {}

private:
    Graph& _g;
    EdgeIndexMap _edge_index;
    CorrProb _corr_prob;
    BlockDeg _blockdeg;

    typedef std::unordered_map<pair<deg_t, deg_t>, double> prob_map_t;
    prob_map_t _probs;
};

// general "traditional" stochastic blockmodel
// this version is based on the alias method, and does not keep the degrees fixed
template <class Graph, class EdgeIndexMap, class CorrProb, class BlockDeg,
          bool micro>
class TradBlockRewireStrategy
{
public:
    typedef typename graph_traits<Graph>::vertex_descriptor vertex_t;
    typedef typename graph_traits<Graph>::edge_descriptor edge_t;
    typedef typename EdgeIndexMap::value_type index_t;
    typedef typename BlockDeg::block_t deg_t;

    TradBlockRewireStrategy(Graph& g, EdgeIndexMap edge_index,
                            vector<edge_t>& edges, CorrProb corr_prob,
                            BlockDeg blockdeg, bool, rng_t& rng,
                            bool parallel_edges, bool configuration)

        : _g(g), _edge_index(edge_index), _edges(edges), _corr_prob(corr_prob),
          _blockdeg(blockdeg), _rng(rng), _sampler(nullptr),
          _configuration(configuration),
          _nmap(get(vertex_index, g), num_vertices(g))
    {
        for (auto v : vertices_range(_g))
        {
            deg_t d = _blockdeg.get_block(v, g);
            _vertices[d].push_back(v);
        }

        if (!micro)
        {
            std::unordered_map<pair<deg_t, deg_t>, double> probs;
            _corr_prob.get_probs(probs);

            vector<double> dprobs;
            if (probs.empty())
            {
                for (auto& s : _vertices)
                {
                    for (auto& t : _vertices)
                    {
                        double p = _corr_prob(s.first, t.first);
                        if (std::isnan(p) || std::isinf(p) || p <= 0)
                            continue;

                        _items.push_back(make_pair(s.first, t.first));
                        dprobs.push_back(p * s.second.size() * t.second.size());
                    }
                }
            }
            else
            {
                for (auto& stp : probs)
                {
                    deg_t s = stp.first.first;
                    deg_t t = stp.first.second;
                    double p = stp.second;
                    // avoid zero probability to not get stuck in rejection step
                    if (std::isnan(p) || std::isinf(p) || p <= 0)
                        continue;
                    _items.push_back(make_pair(s, t));
                    dprobs.push_back(p * _vertices[s].size() * _vertices[t].size());
                }
            }

            if (_items.empty())
                throw GraphException("No connection probabilities larger than zero!");

            _sampler = new Sampler<pair<deg_t, deg_t> >(_items, dprobs);
        }

        if (!configuration || !parallel_edges)
        {
            for (size_t i = 0; i < edges.size(); ++i)
                add_count(source(edges[i], g), target(edges[i], g), _nmap, g);
        }
    }

    ~TradBlockRewireStrategy()
    {
        if (_sampler != nullptr)
            delete _sampler;
    }

    bool operator()(size_t ei, bool self_loops, bool parallel_edges)
    {
        size_t e_s = source(_edges[ei], _g);
        size_t e_t = target(_edges[ei], _g);

        typename graph_traits<Graph>::vertex_descriptor s, t;

        pair<deg_t, deg_t> deg;
        if (micro)
            deg = {_blockdeg.get_block(e_s, _g),
                   _blockdeg.get_block(e_t, _g)};

        while (true)
        {
            if (!micro)
                deg = _sampler->sample(_rng);

            vector<vertex_t>& svs = _vertices[deg.first];
            vector<vertex_t>& tvs = _vertices[deg.second];

            if (svs.empty() || tvs.empty())
                continue;

            s = uniform_sample(svs, _rng);
            t = uniform_sample(tvs, _rng);

            if (!graph_tool::is_directed(_g) &&
                deg.first == deg.second && s != t && self_loops)
            {
                // sample self-loops w/ correct probability for undirected
                // graphs
                std::bernoulli_distribution reject(.5);
                if (reject(_rng))
                    continue;
            }

            break;
        }

        // reject self-loops if not allowed
        if (!self_loops && s == t)
            return false;

        // reject parallel edges if not allowed
        if (!parallel_edges && get_count(s, t, _nmap, _g))
            return false;

        if (!_configuration)
        {
            size_t m = get_count(s, t, _nmap, _g);
            size_t m_e = get_count(e_s, e_t, _nmap, _g);

            double a = (m + 1) / double(m_e);

            std::bernoulli_distribution accept(std::min(a, 1.));
            if (!accept(_rng))
                return false;
        }

        remove_edge(_edges[ei], _g);
        edge_t ne = add_edge(s, t, _g).first;
        _edges[ei] = ne;

        if (!_configuration || !parallel_edges)
        {
            remove_count(e_s, e_t, _nmap, _g);
            add_count(s, t, _nmap, _g);
        }

        return true;
    }

private:
    Graph& _g;
    EdgeIndexMap _edge_index;
    vector<edge_t>& _edges;
    CorrProb _corr_prob;
    BlockDeg _blockdeg;
    rng_t& _rng;

    std::unordered_map<deg_t, vector<vertex_t>> _vertices;

    vector<pair<deg_t, deg_t> > _items;
    Sampler<pair<deg_t, deg_t> >* _sampler;

    bool _configuration;

    typedef gt_hash_map<size_t, size_t> nmapv_t;
    typedef typename property_map_type::apply<nmapv_t,
                                              typename property_map<Graph, vertex_index_t>::type>
        ::type::unchecked_t nmap_t;
    nmap_t _nmap;
};

template <class Graph, class EdgeIndexMap, class CorrProb, class BlockDeg>
using CanTradBlockRewireStrategy =
    TradBlockRewireStrategy<Graph, EdgeIndexMap, CorrProb, BlockDeg, false>;

template <class Graph, class EdgeIndexMap, class CorrProb, class BlockDeg>
using MicroTradBlockRewireStrategy =
    TradBlockRewireStrategy<Graph, EdgeIndexMap, CorrProb, BlockDeg, true>;

} // graph_tool namespace

#endif // GRAPH_REWIRING_HH


//  (C) Copyright John Maddock 2006.
//  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)

#ifndef BOOST_MATH_TOOLS_RATIONAL_HPP
#define BOOST_MATH_TOOLS_RATIONAL_HPP

#ifdef _MSC_VER
#pragma once
#endif

#include <boost/array.hpp>
#include <boost/math/tools/config.hpp>
#include <boost/mpl/int.hpp>

#if BOOST_MATH_POLY_METHOD == 1
#  define BOOST_HEADER() <BOOST_JOIN(boost/math/tools/detail/polynomial_horner1_, BOOST_MATH_MAX_POLY_ORDER).hpp>
#  include BOOST_HEADER()
#  undef BOOST_HEADER
#elif BOOST_MATH_POLY_METHOD == 2
#  define BOOST_HEADER() <BOOST_JOIN(boost/math/tools/detail/polynomial_horner2_, BOOST_MATH_MAX_POLY_ORDER).hpp>
#  include BOOST_HEADER()
#  undef BOOST_HEADER
#elif BOOST_MATH_POLY_METHOD == 3
#  define BOOST_HEADER() <BOOST_JOIN(boost/math/tools/detail/polynomial_horner3_, BOOST_MATH_MAX_POLY_ORDER).hpp>
#  include BOOST_HEADER()
#  undef BOOST_HEADER
#endif
#if BOOST_MATH_RATIONAL_METHOD == 1
#  define BOOST_HEADER() <BOOST_JOIN(boost/math/tools/detail/rational_horner1_, BOOST_MATH_MAX_POLY_ORDER).hpp>
#  include BOOST_HEADER()
#  undef BOOST_HEADER
#elif BOOST_MATH_RATIONAL_METHOD == 2
#  define BOOST_HEADER() <BOOST_JOIN(boost/math/tools/detail/rational_horner2_, BOOST_MATH_MAX_POLY_ORDER).hpp>
#  include BOOST_HEADER()
#  undef BOOST_HEADER
#elif BOOST_MATH_RATIONAL_METHOD == 3
#  define BOOST_HEADER() <BOOST_JOIN(boost/math/tools/detail/rational_horner3_, BOOST_MATH_MAX_POLY_ORDER).hpp>
#  include BOOST_HEADER()
#  undef BOOST_HEADER
#endif

#if 0
//
// This just allows dependency trackers to find the headers
// used in the above PP-magic.
//
#include <boost/math/tools/detail/polynomial_horner1_2.hpp>
#include <boost/math/tools/detail/polynomial_horner1_3.hpp>
#include <boost/math/tools/detail/polynomial_horner1_4.hpp>
#include <boost/math/tools/detail/polynomial_horner1_5.hpp>
#include <boost/math/tools/detail/polynomial_horner1_6.hpp>
#include <boost/math/tools/detail/polynomial_horner1_7.hpp>
#include <boost/math/tools/detail/polynomial_horner1_8.hpp>
#include <boost/math/tools/detail/polynomial_horner1_9.hpp>
#include <boost/math/tools/detail/polynomial_horner1_10.hpp>
#include <boost/math/tools/detail/polynomial_horner1_11.hpp>
#include <boost/math/tools/detail/polynomial_horner1_12.hpp>
#include <boost/math/tools/detail/polynomial_horner1_13.hpp>
#include <boost/math/tools/detail/polynomial_horner1_14.hpp>
#include <boost/math/tools/detail/polynomial_horner1_15.hpp>
#include <boost/math/tools/detail/polynomial_horner1_16.hpp>
#include <boost/math/tools/detail/polynomial_horner1_17.hpp>
#include <boost/math/tools/detail/polynomial_horner1_18.hpp>
#include <boost/math/tools/detail/polynomial_horner1_19.hpp>
#include <boost/math/tools/detail/polynomial_horner1_20.hpp>
#include <boost/math/tools/detail/polynomial_horner2_2.hpp>
#include <boost/math/tools/detail/polynomial_horner2_3.hpp>
#include <boost/math/tools/detail/polynomial_horner2_4.hpp>
#include <boost/math/tools/detail/polynomial_horner2_5.hpp>
#include <boost/math/tools/detail/polynomial_horner2_6.hpp>
#include <boost/math/tools/detail/polynomial_horner2_7.hpp>
#include <boost/math/tools/detail/polynomial_horner2_8.hpp>
#include <boost/math/tools/detail/polynomial_horner2_9.hpp>
#include <boost/math/tools/detail/polynomial_horner2_10.hpp>
#include <boost/math/tools/detail/polynomial_horner2_11.hpp>
#include <boost/math/tools/detail/polynomial_horner2_12.hpp>
#include <boost/math/tools/detail/polynomial_horner2_13.hpp>
#include <boost/math/tools/detail/polynomial_horner2_14.hpp>
#include <boost/math/tools/detail/polynomial_horner2_15.hpp>
#include <boost/math/tools/detail/polynomial_horner2_16.hpp>
#include <boost/math/tools/detail/polynomial_horner2_17.hpp>
#include <boost/math/tools/detail/polynomial_horner2_18.hpp>
#include <boost/math/tools/detail/polynomial_horner2_19.hpp>
#include <boost/math/tools/detail/polynomial_horner2_20.hpp>
#include <boost/math/tools/detail/polynomial_horner3_2.hpp>
#include <boost/math/tools/detail/polynomial_horner3_3.hpp>
#include <boost/math/tools/detail/polynomial_horner3_4.hpp>
#include <boost/math/tools/detail/polynomial_horner3_5.hpp>
#include <boost/math/tools/detail/polynomial_horner3_6.hpp>
#include <boost/math/tools/detail/polynomial_horner3_7.hpp>
#include <boost/math/tools/detail/polynomial_horner3_8.hpp>
#include <boost/math/tools/detail/polynomial_horner3_9.hpp>
#include <boost/math/tools/detail/polynomial_horner3_10.hpp>
#include <boost/math/tools/detail/polynomial_horner3_11.hpp>
#include <boost/math/tools/detail/polynomial_horner3_12.hpp>
#include <boost/math/tools/detail/polynomial_horner3_13.hpp>
#include <boost/math/tools/detail/polynomial_horner3_14.hpp>
#include <boost/math/tools/detail/polynomial_horner3_15.hpp>
#include <boost/math/tools/detail/polynomial_horner3_16.hpp>
#include <boost/math/tools/detail/polynomial_horner3_17.hpp>
#include <boost/math/tools/detail/polynomial_horner3_18.hpp>
#include <boost/math/tools/detail/polynomial_horner3_19.hpp>
#include <boost/math/tools/detail/polynomial_horner3_20.hpp>
#include <boost/math/tools/detail/rational_horner1_2.hpp>
#include <boost/math/tools/detail/rational_horner1_3.hpp>
#include <boost/math/tools/detail/rational_horner1_4.hpp>
#include <boost/math/tools/detail/rational_horner1_5.hpp>
#include <boost/math/tools/detail/rational_horner1_6.hpp>
#include <boost/math/tools/detail/rational_horner1_7.hpp>
#include <boost/math/tools/detail/rational_horner1_8.hpp>
#include <boost/math/tools/detail/rational_horner1_9.hpp>
#include <boost/math/tools/detail/rational_horner1_10.hpp>
#include <boost/math/tools/detail/rational_horner1_11.hpp>
#include <boost/math/tools/detail/rational_horner1_12.hpp>
#include <boost/math/tools/detail/rational_horner1_13.hpp>
#include <boost/math/tools/detail/rational_horner1_14.hpp>
#include <boost/math/tools/detail/rational_horner1_15.hpp>
#include <boost/math/tools/detail/rational_horner1_16.hpp>
#include <boost/math/tools/detail/rational_horner1_17.hpp>
#include <boost/math/tools/detail/rational_horner1_18.hpp>
#include <boost/math/tools/detail/rational_horner1_19.hpp>
#include <boost/math/tools/detail/rational_horner1_20.hpp>
#include <boost/math/tools/detail/rational_horner2_2.hpp>
#include <boost/math/tools/detail/rational_horner2_3.hpp>
#include <boost/math/tools/detail/rational_horner2_4.hpp>
#include <boost/math/tools/detail/rational_horner2_5.hpp>
#include <boost/math/tools/detail/rational_horner2_6.hpp>
#include <boost/math/tools/detail/rational_horner2_7.hpp>
#include <boost/math/tools/detail/rational_horner2_8.hpp>
#include <boost/math/tools/detail/rational_horner2_9.hpp>
#include <boost/math/tools/detail/rational_horner2_10.hpp>
#include <boost/math/tools/detail/rational_horner2_11.hpp>
#include <boost/math/tools/detail/rational_horner2_12.hpp>
#include <boost/math/tools/detail/rational_horner2_13.hpp>
#include <boost/math/tools/detail/rational_horner2_14.hpp>
#include <boost/math/tools/detail/rational_horner2_15.hpp>
#include <boost/math/tools/detail/rational_horner2_16.hpp>
#include <boost/math/tools/detail/rational_horner2_17.hpp>
#include <boost/math/tools/detail/rational_horner2_18.hpp>
#include <boost/math/tools/detail/rational_horner2_19.hpp>
#include <boost/math/tools/detail/rational_horner2_20.hpp>
#include <boost/math/tools/detail/rational_horner3_2.hpp>
#include <boost/math/tools/detail/rational_horner3_3.hpp>
#include <boost/math/tools/detail/rational_horner3_4.hpp>
#include <boost/math/tools/detail/rational_horner3_5.hpp>
#include <boost/math/tools/detail/rational_horner3_6.hpp>
#include <boost/math/tools/detail/rational_horner3_7.hpp>
#include <boost/math/tools/detail/rational_horner3_8.hpp>
#include <boost/math/tools/detail/rational_horner3_9.hpp>
#include <boost/math/tools/detail/rational_horner3_10.hpp>
#include <boost/math/tools/detail/rational_horner3_11.hpp>
#include <boost/math/tools/detail/rational_horner3_12.hpp>
#include <boost/math/tools/detail/rational_horner3_13.hpp>
#include <boost/math/tools/detail/rational_horner3_14.hpp>
#include <boost/math/tools/detail/rational_horner3_15.hpp>
#include <boost/math/tools/detail/rational_horner3_16.hpp>
#include <boost/math/tools/detail/rational_horner3_17.hpp>
#include <boost/math/tools/detail/rational_horner3_18.hpp>
#include <boost/math/tools/detail/rational_horner3_19.hpp>
#include <boost/math/tools/detail/rational_horner3_20.hpp>
#endif

namespace boost{ namespace math{ namespace tools{

//
// Forward declaration to keep two phase lookup happy:
//
template <class T, class U>
U evaluate_polynomial(const T* poly, U const& z, std::size_t count);

namespace detail{

template <class T, class V, class Tag>
inline V evaluate_polynomial_c_imp(const T* a, const V& val, const Tag*)
{
   return evaluate_polynomial(a, val, Tag::value);
}

} // namespace detail

//
// Polynomial evaluation with runtime size.
// This requires a for-loop which may be more expensive than
// the loop expanded versions above:
//
template <class T, class U>
inline U evaluate_polynomial(const T* poly, U const& z, std::size_t count)
{
   BOOST_ASSERT(count > 0);
   U sum = static_cast<U>(poly[count - 1]);
   for(int i = static_cast<int>(count) - 2; i >= 0; --i)
   {
      sum *= z;
      sum += static_cast<U>(poly[i]);
   }
   return sum;
}
//
// Compile time sized polynomials, just inline forwarders to the
// implementations above:
//
template <std::size_t N, class T, class V>
inline V evaluate_polynomial(const T(&a)[N], const V& val)
{
   typedef mpl::int_<N> tag_type;
   return detail::evaluate_polynomial_c_imp(static_cast<const T*>(a), val, static_cast<tag_type const*>(0));
}

template <std::size_t N, class T, class V>
inline V evaluate_polynomial(const boost::array<T,N>& a, const V& val)
{
   typedef mpl::int_<N> tag_type;
   return detail::evaluate_polynomial_c_imp(static_cast<const T*>(a.data()), val, static_cast<tag_type const*>(0));
}
//
// Even polynomials are trivial: just square the argument!
//
template <class T, class U>
inline U evaluate_even_polynomial(const T* poly, U z, std::size_t count)
{
   return evaluate_polynomial(poly, z*z, count);
}

template <std::size_t N, class T, class V>
inline V evaluate_even_polynomial(const T(&a)[N], const V& z)
{
   return evaluate_polynomial(a, z*z);
}

template <std::size_t N, class T, class V>
inline V evaluate_even_polynomial(const boost::array<T,N>& a, const V& z)
{
   return evaluate_polynomial(a, z*z);
}
//
// Odd polynomials come next:
//
template <class T, class U>
inline U evaluate_odd_polynomial(const T* poly, U z, std::size_t count)
{
   return poly[0] + z * evaluate_polynomial(poly+1, z*z, count-1);
}

template <std::size_t N, class T, class V>
inline V evaluate_odd_polynomial(const T(&a)[N], const V& z)
{
   typedef mpl::int_<N-1> tag_type;
   return a[0] + z * detail::evaluate_polynomial_c_imp(static_cast<const T*>(a) + 1, z*z, static_cast<tag_type const*>(0));
}

template <std::size_t N, class T, class V>
inline V evaluate_odd_polynomial(const boost::array<T,N>& a, const V& z)
{
   typedef mpl::int_<N-1> tag_type;
   return a[0] + z * detail::evaluate_polynomial_c_imp(static_cast<const T*>(a.data()) + 1, z*z, static_cast<tag_type const*>(0));
}

template <class T, class U, class V>
V evaluate_rational(const T* num, const U* denom, const V& z_, std::size_t count);

namespace detail{

template <class T, class U, class V, class Tag>
inline V evaluate_rational_c_imp(const T* num, const U* denom, const V& z, const Tag*)
{
   return boost::math::tools::evaluate_rational(num, denom, z, Tag::value);
}

}
//
// Rational functions: numerator and denominator must be
// equal in size.  These always have a for-loop and so may be less
// efficient than evaluating a pair of polynomials. However, there
// are some tricks we can use to prevent overflow that might otherwise
// occur in polynomial evaluation, if z is large.  This is important
// in our Lanczos code for example.
//
template <class T, class U, class V>
V evaluate_rational(const T* num, const U* denom, const V& z_, std::size_t count)
{
   V z(z_);
   V s1, s2;
   if(z <= 1)
   {
      s1 = static_cast<V>(num[count-1]);
      s2 = static_cast<V>(denom[count-1]);
      for(int i = (int)count - 2; i >= 0; --i)
      {
         s1 *= z;
         s2 *= z;
         s1 += num[i];
         s2 += denom[i];
      }
   }
   else
   {
      z = 1 / z;
      s1 = static_cast<V>(num[0]);
      s2 = static_cast<V>(denom[0]);
      for(unsigned i = 1; i < count; ++i)
      {
         s1 *= z;
         s2 *= z;
         s1 += num[i];
         s2 += denom[i];
      }
   }
   return s1 / s2;
}

template <std::size_t N, class T, class U, class V>
inline V evaluate_rational(const T(&a)[N], const U(&b)[N], const V& z)
{
   return detail::evaluate_rational_c_imp(a, b, z, static_cast<const mpl::int_<N>*>(0));
}

template <std::size_t N, class T, class U, class V>
inline V evaluate_rational(const boost::array<T,N>& a, const boost::array<U,N>& b, const V& z)
{
   return detail::evaluate_rational_c_imp(a.data(), b.data(), z, static_cast<mpl::int_<N>*>(0));
}

} // namespace tools
} // namespace math
} // namespace boost

#endif // BOOST_MATH_TOOLS_RATIONAL_HPP


// smooth_feedback: Control theory on Lie groups
// https://github.com/pettni/smooth_feedback
//
// Licensed under the MIT License <http://opensource.org/licenses/MIT>.
//
// Copyright (c) 2021 Petter Nilsson
//
// 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.

#ifndef SMOOTH__FEEDBACK__MESH_FUNCTION_HPP_
#define SMOOTH__FEEDBACK__MESH_FUNCTION_HPP_

/**
 * @file
 * @brief Evaluate transform-like functions and derivatives on collocation points.
 */

#include <Eigen/Dense>
#include <Eigen/Sparse>

#include <smooth/diff.hpp>

#include "mesh.hpp"
#include "smooth/feedback/utils/sparse.hpp"

namespace smooth::feedback {

template<uint8_t Deriv>
struct MeshValue;

/**
 * @brief Result of Mesh function.
 *
 * A generic function is a mapping (t0, tf, q, X, U) -> R^M
 *
 * The function variables are
 * NAME     SIZE
 *   t0        1
 *   tf        1
 *    X nx*(N+1)
 *    U     nu*N
 *
 * Let the total number of variables be K = 2 + nq + nx*(N+1) + nu*N.
 */
template<>
struct MeshValue<0>
{
  /// @brief Function value (size M)
  Eigen::VectorXd F;

  /// @brief If set to true correct allocation is assumed, and no allocation is performed.
  bool allocated{false};
};

/**
 * @brief Result and first derivative of Mesh function.
 *
 * @see MeshValue<0>
 */
template<>
struct MeshValue<1> : public MeshValue<0>
{
  /// @brief Size M x numVar
  Eigen::SparseMatrix<double> dF;
};

/**
 * @brief Result, first, and second derivatives w.r.t. collocation variables
 *
 * @see MeshValue<0>, MeshValue<1>
 */
template<>
struct MeshValue<2> : public MeshValue<1>
{
  /// @brief Multipliers (must be set before)
  Eigen::VectorXd lambda;

  /// @brief Size numVar x numVar
  Eigen::SparseMatrix<double> d2F;
};

/**
 * @brief Reset MeshValue to zeros.
 *
 * If sparse matrices are compressed the coefficients are set to zero, otherwise coefficients are
 * removed but memory is kept.
 */
template<uint8_t Deriv>
void set_zero(MeshValue<Deriv> & mv)
{
  mv.F.setZero();
  if constexpr (Deriv >= 1) { set_zero(mv.dF); }
  if constexpr (Deriv >= 2) { set_zero(mv.d2F); }
}

/**
 * @brief Evaluate function over a mesh.
 *
 * @tparam Deriv differentiation order (0, 1, or 2)
 * @tparam DT inner differentiation method (used to differentiate f)
 *
 * Computes the expression
 * \f[
 *  \begin{bmatrix}
 *    f(t_0 + (t_f - t_0) \tau_0, x_0, u_0) \\
 *    f(t_0 + (t_f - t_0) \tau_1, x_1, u_1) \\
 *    \vdots \\
 *    f(t_{N-1} + (t_f - t_0) \tau_{N-1}, x_{N-1}, u_{N-1})
 *  \end{bmatrix}.
 * \f]
 *
 * If DT > 0 derivatives w.r.t. t0, tf, {xi} and {ui} are returned.
 *
 * @param out result structure
 * @param m mesh
 * @param f integrand
 * @param t0 initial time parameter
 * @param tf final time parameter
 * @param xs state parameters {xi}
 * @param us input parameters {ui}
 * @param scale scale values with quadrature weights
 */
template<uint8_t Deriv, diff::Type DT = diff::Type::Default>
  requires(Deriv <= 2)
void mesh_eval(
  MeshValue<Deriv> & out,
  const MeshType auto & m,
  auto & f,
  const double t0,
  const double tf,
  std::ranges::range auto && xs,
  std::ranges::range auto && us,
  bool scale = false)
{
  using utils::zip, std::views::iota;
  using X = PlainObject<std::decay_t<std::ranges::range_value_t<decltype(xs)>>>;
  using U = PlainObject<std::decay_t<std::ranges::range_value_t<decltype(us)>>>;

  static constexpr auto nx = Dof<X>;
  static constexpr auto nu = Dof<U>;
  static constexpr auto nf = Dof<std::invoke_result_t<decltype(f), double, X, U>>;

  static_assert(nx > -1, "State dimension must be static");
  static_assert(nu > -1, "Input dimension must be static");
  static_assert(nf > -1, "Output size must be static");

  const auto N = m.N_colloc();

  const Eigen::Index numOuts = nf * N;
  const Eigen::Index numVars = 2 + nx * (N + 1) + nu * N;

  // ALLOCATION

  if (!out.allocated) {
    out.F.resize(numOuts);

    if constexpr (Deriv >= 1) {
      Eigen::VectorXi pattern = Eigen::VectorXi::Zero(numVars);
      pattern.segment(0, 1).setConstant(numOuts);                 // t0 is dense
      pattern.segment(1, 1).setConstant(numOuts);                 // tf is dense
      pattern.segment(2, nx * N).setConstant(nf);                 // block diagonal, last x not used
      pattern.segment(2 + nx * (N + 1), nu * N).setConstant(nf);  // block diagonal

      out.dF.resize(numOuts, numVars);
      out.dF.reserve(pattern);
    }

    if constexpr (Deriv >= 2u) {
      Eigen::VectorXi pattern = Eigen::VectorXi::Zero(numVars);
      pattern.segment(0, 1).setConstant(1);  // t0
      pattern.segment(1, 1).setConstant(2);  // t0, tf
      for (auto i = 0u; i < N; ++i) {
        for (auto j = 0u; j < nx; ++j) {
          pattern(2 + nx * i + j) = 2 + (j + 1);  // t0, tf, x upper diag
        }
      }
      for (auto i = 0u; i < N; ++i) {
        for (auto j = 0u; j < nu; ++j) {
          pattern(2 + nx * (N + 1) + nu * i + j) = 2 + nx + (j + 1);  // t0, tf, x, u upper diag
        }
      }
      out.d2F.resize(numVars, numVars);
      out.d2F.reserve(pattern);
    }

    out.allocated = true;
  }

  set_zero(out);

  for (const auto & [i, tau, w_quad, x, u] :
       zip(iota(0u, N), m.all_nodes(), m.all_weights(), xs, us)) {
    const double ti = t0 + (tf - t0) * tau;
    const X xi      = x;
    const U ui      = u;

    const double mtau = 1. - tau;
    const double w    = scale ? w_quad : 1.;

    const auto fvals = diff::dr<Deriv, DT>(f, wrt(ti, xi, ui));

    out.F.segment(i * nf, nf) = w * std::get<0>(fvals);

    if constexpr (Deriv >= 1u) {
      const auto & df = std::get<1>(fvals);
      const auto row0 = nf * i;

      block_add(out.dF, row0, 0, df.middleCols(0, 1), w * mtau);
      block_add(out.dF, row0, 1, df.middleCols(0, 1), w * tau);
      block_add(out.dF, row0, 2 + i * nx, df.middleCols(1, nx), w);
      block_add(out.dF, row0, 2 + nx * (N + 1) + nu * i, df.middleCols(1 + nx, nu), w);

      if constexpr (Deriv >= 2u) {
        assert(out.lambda.size() == numOuts);

        const auto & d2f = std::get<2>(fvals);

        for (auto j = 0u; j < nf; ++j) {
          const double wl = w * out.lambda(row0 + j);

          // destination block locations
          const auto t0_d = 0;
          const auto tf_d = 1;
          const auto x_d  = 2 + i * nx;
          const auto u_d  = 2 + (N + 1) * nx + i * nu;

          // source block locations
          const auto b_s = (1 + nx + nu) * j;
          const auto t_s = 0;
          const auto x_s = 1;
          const auto u_s = 1 + nx;

          // clang-format off
          // t0 row
          block_add(out.d2F, t0_d, t0_d, d2f.block(t_s, b_s + t_s, 1, 1 ), wl * mtau * mtau, true);
          block_add(out.d2F, t0_d, tf_d, d2f.block(t_s, b_s + t_s, 1, 1 ), wl * mtau * tau, true);
          block_add(out.d2F, t0_d,  x_d, d2f.block(t_s, b_s + x_s, 1, nx), wl * mtau, true);
          block_add(out.d2F, t0_d,  u_d, d2f.block(t_s, b_s + u_s, 1, nu), wl * mtau, true);

          // tf row
          block_add(out.d2F, tf_d, tf_d, d2f.block(t_s, b_s + t_s, 1, 1 ), wl * tau * tau, true);
          block_add(out.d2F, tf_d,  x_d, d2f.block(t_s, b_s + x_s, 1, nx), wl * tau, true);
          block_add(out.d2F, tf_d,  u_d, d2f.block(t_s, b_s + u_s, 1, nu), wl * tau, true);

          // x row
          block_add(out.d2F, x_d,  x_d, d2f.block(x_s, b_s + x_s, nx, nx), wl, true);
          block_add(out.d2F, x_d,  u_d, d2f.block(x_s, b_s + u_s, nx, nu), wl, true);

          // u row
          block_add(out.d2F, u_d,  u_d, d2f.block(u_s, b_s + u_s, nu, nu), wl, true);
          // clang-format on
        }
      }
    }
  }
}

/**
 * @brief Evaluate integral over a mesh.
 *
 * @tparam Deriv differentiation order (0, 1, or 2)
 * @tparam DT inner differentiation method (used to differentiate f)

 * This function approximates the integral
 * \f[
 *   \int_{t_0}^{t_f} f(s, x(s), u(s)) \mathrm{d} s.
 * \f]
 * by computing the quadrature
 * \f[
 *    (t_f - t_0) \sum_{i = 0}^N w_i f(t_0 + (t_f - t_0) \tau_i, x_i, u_i).
 * \f]
 *
 * If DT > 0 derivatives w.r.t. t0, tf, {xi} and {ui} are returned.
 *
 * @param out result structure
 * @param m mesh
 * @param f integrand
 * @param t0 initial time parameter
 * @param tf final time parameter
 * @param xs state parameters {xi}
 * @param us input parameters {ui}
 */
template<uint8_t Deriv, diff::Type DT = diff::Type::Default>
  requires(Deriv <= 2)
void mesh_integrate(
  MeshValue<Deriv> & out,
  const MeshType auto & m,
  auto & f,
  const double t0,
  const double tf,
  std::ranges::range auto && xs,
  std::ranges::range auto && us)
{
  using utils::zip, std::views::iota;
  using X = PlainObject<std::decay_t<std::ranges::range_value_t<decltype(xs)>>>;
  using U = PlainObject<std::decay_t<std::ranges::range_value_t<decltype(us)>>>;

  static constexpr auto nx = Dof<X>;
  static constexpr auto nu = Dof<U>;
  static constexpr auto nf = Dof<std::invoke_result_t<decltype(f), double, X, U>>;

  static_assert(nx > -1, "State dimension must be static");
  static_assert(nu > -1, "Input dimension must be static");
  static_assert(nf > -1, "Output size must be static");

  const auto N = m.N_colloc();

  const Eigen::Index numOuts = nf;
  const Eigen::Index numVars = 2 + nx * (N + 1) + nu * N;

  // ALLOCATION

  if (!out.allocated) {
    out.F.resize(numOuts);

    if constexpr (Deriv >= 1) {
      Eigen::VectorXi pattern = Eigen::VectorXi::Constant(numVars, numOuts);  // dense
      pattern.segment(2 + nx * N, nx).setZero();

      out.dF.resize(numOuts, numVars);
      out.dF.reserve(pattern);
    }

    if constexpr (Deriv >= 2) {
      Eigen::VectorXi pattern = Eigen::VectorXi::Zero(numVars);  // dense
      pattern(0)              = 1;                               // t0 dense
      pattern(1)              = 2;                               // tf dense
      for (auto i = 0u; i < N; ++i) {
        for (auto j = 0u; j < nx; ++j) {
          pattern(2 + nx * i + j) = 2 + (j + 1);  // t0, tf, x upper diag
        }
      }
      for (auto i = 0u; i < N; ++i) {
        for (auto j = 0u; j < nu; ++j) {
          pattern(2 + nx * (N + 1) + nu * i + j) = 2 + nx + (j + 1);  // t0, tf, x, u upper diag
        }
      }

      out.d2F.resize(numVars, numVars);
      out.d2F.reserve(pattern.replicate(numOuts, 1).reshaped());
    }

    out.allocated = true;
  }

  set_zero(out);

  for (const auto & [i, tau, w, x, u] : zip(iota(0u, N), m.all_nodes(), m.all_weights(), xs, us)) {
    const double ti = t0 + (tf - t0) * tau;
    const X xi      = x;
    const U ui      = u;

    const double mtau = 1. - tau;

    const auto fvals = diff::dr<Deriv, DT>(f, wrt(ti, xi, ui));

    const auto & fval = std::get<0>(fvals);
    out.F.noalias() += w * (tf - t0) * fval;

    if constexpr (Deriv >= 1u) {
      const auto & df = std::get<1>(fvals);
      // t0
      block_add(out.dF, 0, 0, df.middleCols(0, 1), w * (tf - t0) * mtau);
      block_add(out.dF, 0, 0, fval, -w);
      // tf
      block_add(out.dF, 0, 1, df.middleCols(0, 1), w * (tf - t0) * tau);
      block_add(out.dF, 0, 1, fval, w);
      // x
      block_add(out.dF, 0, 2 + i * nx, df.middleCols(1, nx), w * (tf - t0));
      // u
      block_add(out.dF, 0, 2 + nx * (N + 1) + nu * i, df.middleCols(1 + nx, nu), w * (tf - t0));

      if constexpr (Deriv >= 2u) {
        assert(out.lambda.size() == numOuts);

        const auto & d2f = std::get<2>(fvals);

        for (auto j = 0u; j < nf; ++j) {
          const double wl = w * out.lambda(j);

          // source block locations
          const auto b_s = (1 + nx + nu) * j;  // horizontal block
          const auto t_s = 0;                  // t
          const auto x_s = 1;                  // x
          const auto u_s = 1 + nx;             // u

          // destination block locations
          const auto t0_d = 0;                          // t0
          const auto tf_d = 1;                          // tf
          const auto x_d  = 2 + nx * i;                 // x[i]
          const auto u_d  = 2 + nx * (N + 1) + nu * i;  // u[i]

          // clang-format off
          // t0t0
          block_add(out.d2F, t0_d, t0_d, d2f.block(t_s, b_s + t_s, 1,  1), wl * (tf - t0) * mtau * mtau, true);
          block_add(out.d2F, t0_d, t0_d,  df.block(j,         t_s, 1,  1), -wl * 2 * mtau, true);
          // t0tf
          block_add(out.d2F, t0_d, tf_d, d2f.block(t_s, b_s + t_s, 1,  1), wl * (tf - t0) * mtau * tau, true);
          block_add(out.d2F, t0_d, tf_d,  df.block(j,         t_s, 1,  1), wl * (1 - 2 * tau), true);
          // t0x
          block_add(out.d2F, t0_d, x_d,  d2f.block(t_s, b_s + x_s, 1, nx), wl * (tf - t0) * mtau, true);
          block_add(out.d2F, t0_d, x_d,   df.block(j,         x_s, 1, nx), -wl, true);
          // t0u
          block_add(out.d2F, t0_d, u_d,  d2f.block(t_s, b_s + u_s, 1, nu), wl * (tf - t0) * mtau, true);
          block_add(out.d2F, t0_d, u_d,   df.block(j,         u_s, 1, nu), -wl, true);

          // tftf
          block_add(out.d2F, tf_d, tf_d, d2f.block(t_s, b_s + t_s, 1,  1), wl * (tf - t0) * tau * tau, true);
          block_add(out.d2F, tf_d, tf_d,  df.block(j,         t_s, 1,  1), wl * 2 * tau, true);
          // tfx
          block_add(out.d2F, tf_d, x_d,  d2f.block(t_s, b_s + x_s, 1, nx), wl * (tf - t0) * tau, true);
          block_add(out.d2F, tf_d, x_d,   df.block(j,         x_s, 1, nx), wl, true);
          // tfu
          block_add(out.d2F, tf_d, u_d,  d2f.block(t_s, b_s + u_s, 1, nu), wl * (tf - t0) * tau, true);
          block_add(out.d2F, tf_d, u_d,   df.block(j,         u_s, 1, nu), wl, true);

          // xx
          block_add(out.d2F, x_d, x_d,  d2f.block(x_s, b_s + x_s, nx, nx), wl * (tf - t0), true);
          // xu
          block_add(out.d2F, x_d, u_d,  d2f.block(x_s, b_s + u_s, nx, nu), wl * (tf - t0), true);

          // uu
          block_add(out.d2F, u_d, u_d,  d2f.block(u_s, b_s + u_s, nu, nu), wl * (tf - t0), true);
          // clang-format on
        }
      }
    }
  }
}

/**
 * @brief Evaluate dynamic constraints over a mesh (differentiation version).
 *
 * @tparam Deriv differentiation order (0, 1, or 2)
 * @tparam DT inner differentiation method (used to differentiate f)
 *
 * For interval \f$j\f$ with nodes \f$\tau_i\f$ and weights \f$w_i\f$ the corresponding dynamic
 * constraints are a block of the form
 * \f[
 *   \begin{bmatrix}
 *     w_0 \left( f\left(t_0 + (t_f - t_0) \tau_0, x_0, u_0\right) - \sum_{k=0}^{N_j} D_{k, 0} x_k
 * \right) \\
 *     \vdots  \\
 *     w_{N_j-1} \left( f\left(t_0 + (t_f - t_0) \tau_{N_j-1}, x_{N_j - 1}, u_{N_j - 1}\right) -
 * \sum_{k=0}^{N_j} D_{k, N_j-1} x_k \right) \end{bmatrix}, \f] where \f$D\f$ is the interval
 * differentiation matrix (see interval_diffmat() in Mesh).
 *
 * This function returnes all such blocks stacked into a 1D vector.
 *
 * If DT > 0 derivatives w.r.t. t0, tf, {xi} and {ui} are returned.
 *
 * @param out result structure
 * @param m mesh
 * @param f right-hand side of dynamics
 * @param t0 initial time parameter
 * @param tf final time parameter
 * @param xs state parameters {xi}
 * @param us input parameters {ui}
 */
template<uint8_t Deriv, diff::Type DT = diff::Type::Default>
  requires(Deriv <= 2)
void mesh_dyn(
  MeshValue<Deriv> & out,
  const MeshType auto & m,
  auto & f,
  const double t0,
  const double tf,
  std::ranges::range auto && xs,
  std::ranges::range auto && us)
{
  using utils::zip, std::views::iota;
  using X = PlainObject<std::decay_t<std::ranges::range_value_t<decltype(xs)>>>;
  using U = PlainObject<std::decay_t<std::ranges::range_value_t<decltype(us)>>>;

  static constexpr auto nx = Dof<X>;
  static constexpr auto nu = Dof<U>;
  static constexpr auto nf = Dof<std::invoke_result_t<decltype(f), double, X, U>>;

  static_assert(nx > -1, "State dimension must be static");
  static_assert(nu > -1, "Input dimension must be static");
  static_assert(nf > -1, "Output size must be static");
  static_assert(nx == nf, "Output dimension must be same as state dimension");

  const auto N               = m.N_colloc();
  const Eigen::Index numOuts = nx * N;
  const Eigen::Index numVars = 2 + nx * (N + 1) + nu * N;

  // ALLOCATION

  if (!out.allocated) {
    out.F.resize(numOuts);

    if constexpr (Deriv >= 1) {
      Eigen::VectorXi pattern = Eigen::VectorXi::Zero(numVars);
      pattern(0)              = numOuts;  // t0 dense
      pattern(1)              = numOuts;  // tf dense

      // x has blocks depending on mesh size
      auto idx0 = 2u;
      for (auto ival = 0u; ival < m.N_ivals(); ++ival) {
        const std::size_t K = m.N_colloc_ival(ival);
        pattern.segment(idx0, K * nx) += Eigen::VectorXi::Constant(K * nx, nx + K - 1);
        pattern.segment(idx0 + K * nx, nx) += Eigen::VectorXi::Constant(nx, K);
        idx0 += K * nx;
      }

      // u is block diagonal with small blocks
      pattern.segment(2 + nx * (N + 1), nu * N).setConstant(nx);

      out.dF.resize(numOuts, numVars);
      out.dF.reserve(pattern);
    }

    if constexpr (Deriv >= 2) {
      Eigen::VectorXi pattern = Eigen::VectorXi::Zero(numVars);
      pattern(0)              = 1;
      pattern(1)              = 2;
      for (auto i = 0u; i < N; ++i) {
        for (auto j = 0u; j < nx; ++j) {
          pattern(2 + nx * i + j) = 2 + (j + 1);  // t0, tf, x upper diag
        }
      }
      for (auto i = 0u; i < N; ++i) {
        for (auto j = 0u; j < nu; ++j) {
          pattern(2 + nx * (N + 1) + nu * i + j) = 2 + nx + (j + 1);  // t0, tf, x, u upper diag
        }
      }

      out.d2F.resize(numVars, numVars);
      out.d2F.reserve(pattern);
    }

    out.allocated = true;
  }

  set_zero(out);

  // We build the constraint through two loops over xi:
  //   - the first loop adds wk * f(tk, xk, uk)
  //   - the second loop adds -wk * [x0 ... Xni] * dk

  // ADD FIRST PART

  for (const auto & [i, tau, w, x, u] : zip(iota(0u, N), m.all_nodes(), m.all_weights(), xs, us)) {
    const double ti = t0 + (tf - t0) * tau;
    const X xi      = x;
    const U ui      = u;

    const auto row0   = nx * i;
    const double mtau = 1. - tau;

    const auto fvals = diff::dr<Deriv, DT>(f, wrt(ti, xi, ui));
    const auto fval  = std::get<0>(fvals);

    out.F.segment(row0, nx) += w * (tf - t0) * fval;

    if constexpr (Deriv >= 1) {
      const auto & df = std::get<1>(fvals);

      // clang-format off
      // dF/dt0 = -f + (tf - t0) * df/dti * (1-tau)
      block_add(out.dF, row0, 0, fval, -w);
      block_add(out.dF, row0, 0, df.col(0), w * (tf - t0) * mtau);
      // dF/dtf = f + (tf - t0) * df/dti * tau
      block_add(out.dF, row0, 1, fval, w);
      block_add(out.dF, row0, 1, df.col(0), w * (tf - t0) * tau);
      // dF/dx
      block_add(out.dF, row0, 2 + nx * i, df.middleCols(1, nx), w * (tf - t0));
      // dF/du
      block_add(out.dF, row0, 2 + nx * (N + 1) + nu * i, df.middleCols(1 + nx, nu), w * (tf - t0));
      // clang-format on

      if constexpr (Deriv >= 2) {
        assert(out.lambda.size() == numOuts);

        const auto & d2f = std::get<2>(fvals);

        for (auto j = 0u; j < nx; ++j) {
          const double wl = w * out.lambda(row0 + j);

          // destination block locations
          const auto t0_d = 0;
          const auto tf_d = 1;
          const auto x_d  = 2 + i * nx;
          const auto u_d  = 2 + (N + 1) * nx + i * nu;

          // source block locations
          const auto b_s = (1 + nx + nu) * j;
          const auto t_s = 0;
          const auto x_s = 1;
          const auto u_s = 1 + nx;

          // clang-format off
          // t0t0
          block_add(out.d2F, t0_d, t0_d, d2f.block(t_s, b_s + t_s, 1, 1 ), wl * (tf - t0) * mtau * mtau, true);
          block_add(out.d2F, t0_d, t0_d,  df.block(j,         t_s, 1, 1 ), -wl * 2 * mtau, true);
          // t0tf
          block_add(out.d2F, t0_d, tf_d, d2f.block(t_s, b_s + t_s, 1, 1 ), wl * (tf - t0) * mtau * tau, true);
          block_add(out.d2F, t0_d, tf_d,  df.block(j,         t_s, 1, 1 ), wl * (1. - 2 * tau), true);
          // t0x
          block_add(out.d2F, t0_d, x_d, d2f.block(t_s, b_s + x_s, 1, nx), wl * (tf - t0) * mtau, true);
          block_add(out.d2F, t0_d, x_d,  df.block(j,         x_s, 1, nx), -wl, true);
          // t0u
          block_add(out.d2F, t0_d, u_d, d2f.block(t_s, b_s + u_s, 1, nu), wl * (tf - t0) * mtau, true);
          block_add(out.d2F, t0_d, u_d,  df.block(j,         u_s, 1, nu), -wl, true);

          // tftf
          block_add(out.d2F, tf_d, tf_d, d2f.block(t_s, b_s + t_s, 1,  1), wl * (tf - t0) * tau * tau, true);
          block_add(out.d2F, tf_d, tf_d,  df.block(j,         t_s, 1,  1), wl * 2 * tau, true);
          // tfx
          block_add(out.d2F, tf_d, x_d,  d2f.block(t_s, b_s + x_s, 1, nx), wl * (tf - t0) * tau, true);
          block_add(out.d2F, tf_d, x_d,   df.block(j,         x_s, 1, nx), wl, true);
          // tfu
          block_add(out.d2F, tf_d, u_d,  d2f.block(t_s, b_s + u_s, 1, nu), wl * (tf - t0) * tau, true);
          block_add(out.d2F, tf_d, u_d,   df.block(j,         u_s, 1, nu), wl, true);

          // xx
          block_add(out.d2F, x_d, x_d,  d2f.block(x_s, b_s + x_s, nx, nx), wl * (tf - t0), true);
          // xu
          block_add(out.d2F, x_d, u_d,  d2f.block(x_s, b_s + u_s, nx, nu), wl * (tf - t0), true);

          // uu
          block_add(out.d2F, u_d, u_d,  d2f.block(u_s, b_s + u_s, nu, nu), wl * (tf - t0), true);
          // clang-format on
        }
      }
    }
  }

  // ADD SECOND PART (LINEAR IN X, NO SECOND DERIVATIVE..)

  auto ival      = 0u;  // current interval index
  auto ival_idx0 = 0u;  // node start index of current interval
  auto Nival     = m.N_colloc_ival(ival);

  for (const auto & [i, x] : zip(iota(0u, N + 1), xs)) {
    if (i == ival_idx0 + Nival) {
      // jumping to new interval --- add overlap to current interval before switching
      const auto [alpha, Dus] = m.interval_diffmat_unscaled(ival);
      for (const auto & [j, w] : zip(iota(0u, Nival), m.interval_weights(ival))) {
        const auto row0   = (ival_idx0 + j) * nx;
        const double coef = -w * alpha * Dus(i - ival_idx0, j);

        out.F.segment(row0, nx) += coef * x;

        if constexpr (Deriv >= 1) {
          // add diagonal matrix
          for (auto k = 0u; k < nx; ++k) { out.dF.coeffRef(row0 + k, 2 + nx * i + k) += coef; }
        }
      }

      // update interval
      ++ival;
      if (ival < m.N_ivals()) {
        ival_idx0 += Nival;
        Nival = m.N_colloc_ival(ival);
      }
    }

    if (i < N) {
      const auto [alpha, Dus] = m.interval_diffmat_unscaled(ival);
      for (const auto & [j, w] : zip(iota(0u, Nival), m.interval_weights(ival))) {
        const auto row0   = (ival_idx0 + j) * nx;
        const double coef = -w * alpha * Dus(i - ival_idx0, j);

        out.F.segment(row0, nx) += coef * x;

        if constexpr (Deriv >= 1) {
          // add diagonal matrix
          for (auto k = 0u; k < nx; ++k) { out.dF.coeffRef(row0 + k, 2 + nx * i + k) += coef; }
        }
      }
    }
  }
}

}  // namespace smooth::feedback

#endif  // SMOOTH__FEEDBACK__MESH_FUNCTION_HPP_


/***********************************************************************************************************************
 *  OpenStudio(R), Copyright (c) 2008-2018, Alliance for Sustainable Energy, LLC. All rights reserved.
 *
 *  Redistribution and use in source and binary forms, with or without modification, are permitted provided that the
 *  following conditions are met:
 *
 *  (1) Redistributions of source code must retain the above copyright notice, this list of conditions and the following
 *  disclaimer.
 *
 *  (2) Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the
 *  following disclaimer in the documentation and/or other materials provided with the distribution.
 *
 *  (3) Neither the name of the copyright holder nor the names of any contributors may be used to endorse or promote
 *  products derived from this software without specific prior written permission from the respective party.
 *
 *  (4) Other than as required in clauses (1) and (2), distributions in any form of modifications or other derivative
 *  works may not use the "OpenStudio" trademark, "OS", "os", or any other confusingly similar designation without
 *  specific prior written permission from Alliance for Sustainable Energy, LLC.
 *
 *  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,
 *  INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 *  DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER, THE UNITED STATES GOVERNMENT, OR ANY CONTRIBUTORS BE LIABLE FOR
 *  ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
 *  PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
 *  AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 *  ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 **********************************************************************************************************************/

#ifndef UTILITIES_GEOMETRY_INTERSECTION_HPP
#define UTILITIES_GEOMETRY_INTERSECTION_HPP

#include "../UtilitiesAPI.hpp"

#include "Point3d.hpp"

#include <vector>
#include <boost/optional.hpp>

namespace openstudio{

  /** IntersectionResult contains detailed information about an intersection. */
  class UTILITIES_API IntersectionResult {
  public:
    IntersectionResult(const std::vector<Point3d>& polygon1,
                       const std::vector<Point3d>& polygon2,
                       const std::vector< std::vector<Point3d> >& newPolygons1,
                       const std::vector< std::vector<Point3d> >& newPolygons2);

    // vertices of first polygon after intersection
    std::vector<Point3d> polygon1() const;

    // vertices of second polygon after intersection
    std::vector<Point3d> polygon2() const;

    // new polygons generated from the first surface
    std::vector< std::vector<Point3d> > newPolygons1() const;

    // new polygons generated from the second surface
    std::vector< std::vector<Point3d> > newPolygons2() const;

  private:
    std::vector<Point3d> m_polygon1;
    std::vector<Point3d> m_polygon2;
    std::vector< std::vector<Point3d> > m_newPolygons1;
    std::vector< std::vector<Point3d> > m_newPolygons2;
  };

  /// removes spikes from a polygon, requires that all vertices are in clockwise order on the z = 0 plane (i.e. in face coordinates but reversed)
  UTILITIES_API std::vector<Point3d> removeSpikes(const std::vector<Point3d>& polygon, double tol);

  /// returns true if point is inside polygon, requires that all vertices are in clockwise order on the z = 0 plane (i.e. in face coordinates but reversed)
  UTILITIES_API bool pointInPolygon(const Point3d& point, const std::vector<Point3d>& polygon, double tol);

  /// compute the union of two overlapping polygons, requires that all vertices are in clockwise order on the z = 0 plane (i.e. in face coordinates but reversed)
  UTILITIES_API boost::optional<std::vector<Point3d> > join(const std::vector<Point3d>& polygon1, const std::vector<Point3d>& polygon2, double tol);

  /// compute the union of many polygons, requires that all vertices are in clockwise order on the z = 0 plane (i.e. in face coordinates but reversed)
  UTILITIES_API std::vector<std::vector<Point3d> > joinAll(const std::vector<std::vector<Point3d> >& polygons, double tol);

  /// intersect two polygons, requires that all vertices are in clockwise order on the z = 0 plane (i.e. in face coordinates but reversed)
  UTILITIES_API boost::optional<IntersectionResult> intersect(const std::vector<Point3d>& polygon1, const std::vector<Point3d>& polygon2, double tol);

  /// subtract all holes from polygon, requires that all vertices are in clockwise order on the z = 0 plane (i.e. in face coordinates but reversed)
  UTILITIES_API std::vector<std::vector<Point3d> > subtract(const std::vector<Point3d>& polygon, const std::vector<std::vector<Point3d> >& holes, double tol);

  /// returns true polygon intersects iteself, requires that all vertices are in clockwise order on the z = 0 plane (i.e. in face coordinates but reversed)
  /// returns false if polygon has less than three vertices
  UTILITIES_API bool selfIntersects(const std::vector<Point3d>& polygon, double tol);

  /// returns true if polygon1 intersects polygon2, requires that all vertices are in clockwise order on the z = 0 plane (i.e. in face coordinates but reversed)
  /// returns false if either polygon has less than three vertices
  UTILITIES_API bool intersects(const std::vector<Point3d>& polygon1, const std::vector<Point3d>& polygon2, double tol);

  /// returns true if geometry1 is completely within polygon2, requires that all vertices are in clockwise order on the z = 0 plane (i.e. in face coordinates but reversed)
  /// geometry1 can be a point or a polygon
  /// currently only tests that all points of geometry1 are within polygon2, better support when upgrade to boost 1.57
  UTILITIES_API bool within(const Point3d& point1, const std::vector<Point3d>& polygon2, double tol);
  UTILITIES_API bool within(const std::vector<Point3d>& geometry1, const std::vector<Point3d>& polygon2, double tol);

  /// simplify a list of vertices
  UTILITIES_API std::vector<Point3d> simplify(const std::vector<Point3d>& vertices, bool removeCollinear, double tol);

} // openstudio

#endif //UTILITIES_GEOMETRY_INTERSECTION_HPP


// Copyright (c) 2017 Jeffrey Ichnowski
// All rights reserved.
//
// BSD 3 Clause
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions
// are met:
// 1. Redistributions of source code must retain the above copyright
//    notice, this list of conditions and the following disclaimer.
// 2. Redistributions in binary form must reproduce the above copyright
//    notice, this list of conditions and the following disclaimer in the
//    documentation and/or other materials provided with the distribution.
// 3. Neither the name of the copyright holder nor the names of its
//    contributors may be used to endorse or promote products derived
//    from this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
// FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
// COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
// INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
// (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
// HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
// STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED
// OF THE POSSIBILITY OF SUCH DAMAGE.
#pragma once
#ifndef UNC_ROBOTICS_KDTREE_SPACES_HPP
#define UNC_ROBOTICS_KDTREE_SPACES_HPP

#include <Eigen/Dense>
#include <tuple>
#include <functional>
#include <utility>
#include <type_traits>
#include <ratio>
#include "_tuple.hpp"

namespace unc {
namespace robotics {
namespace kdtree {

namespace detail {
template <typename _Scalar, int _dimensions>
class L2SpaceBase {
public:
    typedef _Scalar Distance;
    typedef Eigen::Matrix<_Scalar, _dimensions, 1> State;

    template <typename _Derived>
    bool isValid(const Eigen::MatrixBase<_Derived>& q) const {
        return q.rows() == _dimensions && q.cols() == 1 && q.allFinite();
    }

    template <typename _DerivedA, typename _DerivedB>
    constexpr Distance distance(
        const Eigen::MatrixBase<_DerivedA>& a,
        const Eigen::MatrixBase<_DerivedB>& b) const
    {
        return (a - b).norm();
    }

    template <typename _DerivedA, typename _DerivedB>
    constexpr State interpolate(
        const Eigen::MatrixBase<_DerivedA>& from,
        const Eigen::MatrixBase<_DerivedB>& to,
        Distance t) const
    {
        return from + (to - from) * t;
    }
};
}

template <typename _Scalar, int _dimensions>
class L2Space : public detail::L2SpaceBase<_Scalar, _dimensions> {
public:
    L2Space(unsigned dimensions = _dimensions) {
        assert(dimensions == _dimensions);
    }

    constexpr unsigned dimensions() const { return _dimensions; }

    template <typename _Derived>
    bool isValid(const Eigen::MatrixBase<_Derived>& q) const {
        return q.rows() == _dimensions && detail::L2SpaceBase<_Scalar, _dimensions>::isValid(q);
    }
};

template <typename _Scalar>
class L2Space<_Scalar, Eigen::Dynamic> : public detail::L2SpaceBase<_Scalar, Eigen::Dynamic> {
    unsigned dimensions_;

public:
    L2Space(unsigned dimensions)
        : dimensions_(dimensions)
    {
    }

    constexpr unsigned dimensions() const {
        return dimensions_;
    }

    template <typename _Derived>
    bool isValid(const Eigen::MatrixBase<_Derived>& q) const {
        return q.rows() == dimensions_ && detail::L2SpaceBase<_Scalar, Eigen::Dynamic>::isValid(q);
    }
};

template <typename _Scalar, int _dimensions>
class BoundedL2Space : public L2Space<_Scalar, _dimensions> {
    Eigen::Array<_Scalar, _dimensions, 2> bounds_;

    typedef typename Eigen::Array<_Scalar, _dimensions, 2>::Index Index;

    void checkBounds() {
        assert((bounds_.col(0) < bounds_.col(1)).all());
        assert((bounds_.col(1) - bounds_.col(0)).allFinite());
    }

public:
    using typename L2Space<_Scalar, _dimensions>::State;

    template <typename _Derived>
    BoundedL2Space(const Eigen::DenseBase<_Derived>& bounds)
        : L2Space<_Scalar, _dimensions>(bounds.rows()),
          bounds_(bounds)
    {
        checkBounds();
    }

    template <typename _DerivedMin, typename _DerivedMax>
    BoundedL2Space(
        const Eigen::DenseBase<_DerivedMin>& min,
        const Eigen::DenseBase<_DerivedMax>& max)
        : L2Space<_Scalar, _dimensions>(min.rows())
    {
        bounds_.col(0) = min;
        bounds_.col(1) = max;
    }

    template <typename _Derived>
    bool isValid(const Eigen::MatrixBase<_Derived>& q) const {
        return L2Space<_Scalar, _dimensions>::isValid(q)
            && (bounds_.col(0) <= q).all()
            && (bounds_.col(1) >= q).all();
    }

    const Eigen::Array<_Scalar, _dimensions, 2>& bounds() const {
        return bounds_;
    }

    _Scalar bounds(Index dim, Index j) const {
        return bounds_(dim, j);
    }
};

template <typename _Scalar>
class SO3Space {
public:
    typedef _Scalar Distance;
    typedef Eigen::Quaternion<_Scalar> State;

    constexpr unsigned dimensions() const { return 3; }

    template <typename _Derived>
    bool isValid(const Eigen::QuaternionBase<_Derived>& q) const {
        return std::abs(1 - q.coeffs().squaredNorm()) <= 1e-5;
    }

    template <typename _DerivedA, typename _DerivedB>
    inline Distance distance(
        const Eigen::QuaternionBase<_DerivedA>& a,
        const Eigen::QuaternionBase<_DerivedB>& b) const
    {
        Distance dot = std::abs(a.coeffs().matrix().dot(b.coeffs().matrix()));
        return dot < 0 ? M_PI_2 : dot > 1 ? 0 : std::acos(dot);
    }

    template <typename _DerivedA, typename _DerivedB>
    constexpr State interpolate(
        const Eigen::QuaternionBase<_DerivedA>& from,
        const Eigen::QuaternionBase<_DerivedB>& to,
        Distance t) const
    {
        Distance dq = from.coeffs().matrix().dot(to.coeffs().matrix());
        if (std::abs(dq) >= 1)
            return from;

        Distance theta = std::acos(std::abs(dq));
        Distance d = 1 / std::sin(theta);
        Distance s0 = std::sin((1 - t) * theta);
        Distance s1 = std::sin(t * theta);

        if (dq < 0)
            s1 = -s1;

        return State(d * (from.coeffs() * s0 + to.coeffs() * s1));
    }
};

template <typename _Scalar>
class SO3AltSpace {
public:
    typedef _Scalar Distance;
    typedef Eigen::Quaternion<_Scalar> State;

    constexpr unsigned dimensions() const { return 3; }

    template <typename _Derived>
    bool isValid(const Eigen::QuaternionBase<_Derived>& q) const {
        return std::abs(1 - q.coeffs().squaredNorm()) <= 1e-5;
    }

    template <typename _DerivedA, typename _DerivedB>
    inline Distance distance(
        const Eigen::QuaternionBase<_DerivedA>& a,
        const Eigen::QuaternionBase<_DerivedB>& b) const
    {
        Distance dot = std::abs(a.coeffs().matrix().dot(b.coeffs().matrix()));
        return dot < 0 ? M_PI_2 : dot > 1 ? 0 : std::acos(dot);
    }
};

template <typename _Scalar>
class SO3RLSpace {
public:
    typedef _Scalar Distance;
    typedef Eigen::Quaternion<_Scalar> State;

    constexpr unsigned dimensions() const { return 4; }

    template <typename _Derived>
    bool isValid(const Eigen::QuaternionBase<_Derived>& q) const {
        return std::abs(1 - q.coeffs().squaredNorm()) <= 1e-5;
    }

    template <typename _DerivedA, typename _DerivedB>
    inline Distance distance(
        const Eigen::QuaternionBase<_DerivedA>& a,
        const Eigen::QuaternionBase<_DerivedB>& b) const
    {
        Distance dot = std::abs(a.coeffs().matrix().dot(b.coeffs().matrix()));
        return dot < 0 ? M_PI_2 : dot > 1 ? 0 : std::acos(dot);
    }
};

template <typename _Space, typename _Ratio = std::ratio<1>>
class RatioWeightedSpace : public _Space {
public:
    typedef _Ratio Ratio;

    static constexpr std::intmax_t num = Ratio::num;
    static constexpr std::intmax_t den = Ratio::den;

    // inherit constructor
    using _Space::_Space;

    // RatioWeightedSpace() {}

    RatioWeightedSpace(const _Space& space)
        : _Space(space)
    {
    }

    RatioWeightedSpace(_Space&& space)
        : _Space(std::forward<_Space>(space))
    {
    }

    template <typename _A, typename _B>
    inline typename _Space::Distance distance(const _A& a, const _B& b) const {
        return _Space::distance(a, b) * num / den;
    }
};

template <std::intmax_t num, std::intmax_t den = 1, typename _Space>
auto makeRatioWeightedSpace(_Space&& space) {
    return RatioWeightedSpace<_Space, std::ratio<num, den>>(std::forward<_Space>(space));
}

template <typename _Space>
class WeightedSpace : public _Space {
public:
    using typename _Space::Distance;

private:
    Distance weight_;

public:
    WeightedSpace(Distance weight, const _Space& space = _Space())
        : _Space(space), weight_(weight)
    {
    }

    WeightedSpace(Distance weight, _Space&& space)
        : _Space(std::forward<_Space>(space)), weight_(weight)
    {
    }

    template <typename ... _Args>
    WeightedSpace(Distance weight, _Args&& ... args)
        : _Space(std::forward<_Args>(args)...),
          weight_(weight)
    {
    }

    constexpr Distance weight() const {
        return weight_;
    }

    template <typename _A, typename _B>
    constexpr Distance distance(const _A& a, const _B& b) const {
        return _Space::distance(a, b) * weight_;
    }
};

template <typename ... _Spaces>
class CompoundSpace {
    typedef std::tuple<_Spaces...> Spaces;

    Spaces spaces_;

    static_assert(sizeof...(_Spaces) > 1, "compound space must have two or more subspaces");

public:
    typedef std::tuple<typename _Spaces::State...> State;
    typedef typename detail::SumResultType<typename _Spaces::Distance...>::type Distance;

    explicit CompoundSpace() {
    }

    explicit CompoundSpace(const _Spaces& ... args)
        : spaces_(args...)
    {
    }

    template <typename ... _Args>
    explicit CompoundSpace(_Args&& ... args)
        : spaces_(std::forward<_Args>(args)...)
    {
    }

    template <std::size_t I>
    constexpr typename std::tuple_element<I, Spaces>::type& get() {
        return std::get<I>(spaces_);
    }

    template <std::size_t I>
    constexpr typename std::tuple_element<I, Spaces>::type const& get() const {
        return std::get<I>(spaces_);
    }

    inline constexpr unsigned dimensions() const {
        using namespace detail;
        return sum(map([](const auto& space) { return space.dimensions(); }, spaces_));
    }

    template <typename _State>
    bool isValid(_State&& q) const {
        using namespace detail;
        // TODO: return reduce(std::logical_and<bool>(), zip([](auto&& subs, auto&& subq) { return subs.isValid(subq); }, spaces_, q));
        assert(false);
        return false;
    }

    template <typename _StateA, typename _StateB>
    inline Distance distance(_StateA&& a, _StateB&& b) const {
        using namespace detail;
        return sum(zip([](auto&& subs, auto&& suba, auto&& subb) {
                    return subs.distance(suba, subb);
                }, spaces_, a, b));
    }

private:
    template <typename _StateA, typename _StateB, std::size_t ... I>
    constexpr State interpolate(
        const _StateA& from,
        const _StateB& to,
        Distance t,
        std::index_sequence<I...>) const
    {
        return State(std::get<I>(spaces_).interpolate(std::get<I>(from), std::get<I>(to), t)...);
    }

public:
    template <typename _StateA, typename _StateB>
    constexpr State interpolate(
        const _StateA& from,
        const _StateB& to,
        Distance t) const
    {
        return interpolate(from, to, t, std::make_index_sequence<sizeof...(_Spaces)>{});
    }

};

template <typename ... _Spaces>
constexpr auto makeCompoundSpace(_Spaces&&... args) {
    return CompoundSpace<typename std::decay<_Spaces>::type...>(std::forward<_Spaces>(args)...);
}

template <typename _Space>
constexpr auto makeWeightedSpace(
    typename _Space::Distance weight,
    _Space&& space = _Space())
{
    return WeightedSpace<_Space>(weight, std::forward<_Space>(space));
}

template <typename _Scalar, std::intmax_t _qWeight = 1, std::intmax_t _tWeight = 1>
using SE3Space = CompoundSpace<
    RatioWeightedSpace<SO3Space<_Scalar>, std::ratio<_qWeight>>,
    RatioWeightedSpace<L2Space<_Scalar, 3>, std::ratio<_tWeight>>>;

template <typename _Scalar, std::intmax_t _qWeight = 1, std::intmax_t _tWeight = 1>
using BoundedSE3Space = CompoundSpace<
    RatioWeightedSpace<SO3Space<_Scalar>, std::ratio<_qWeight>>,
    RatioWeightedSpace<BoundedL2Space<_Scalar, 3>, std::ratio<_tWeight>>>;

}
}
}

namespace std {

template <std::size_t I, class ... _Spaces>
class tuple_element<I, unc::robotics::kdtree::CompoundSpace<_Spaces...>> {
public:
    typedef typename std::tuple_element<I, std::tuple<_Spaces...>>::type type;
};

template <std::size_t I, class ... _Spaces>
constexpr typename std::tuple_element<I, std::tuple<_Spaces...>>::type&
get(unc::robotics::kdtree::CompoundSpace<_Spaces...>& space) {
    return space.template get<I>();
}

template <std::size_t I, class ... _Spaces>
constexpr typename std::tuple_element<I, std::tuple<_Spaces...>>::type const&
get(const unc::robotics::kdtree::CompoundSpace<_Spaces...>& space) {
    return space.template get<I>();
}

}

#endif // UNC_ROBOTICS_KDTREE_SPACES_HPP