//=======================================================================
// Copyright 2009 Trustees of Indiana University.
// Authors: Michael Hansen, Andrew Lumsdaine
//
// Distributed under the Boost Software License, Version 1.0. (See
// accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt)
//=======================================================================

#ifndef BOOST_GRAPH_MCGREGOR_COMMON_SUBGRAPHS_HPP
#define BOOST_GRAPH_MCGREGOR_COMMON_SUBGRAPHS_HPP

#include <algorithm>
#include <vector>
#include <stack>

#include <boost/make_shared.hpp>
#include <boost/graph/adjacency_list.hpp>
#include <boost/graph/filtered_graph.hpp>
#include <boost/graph/graph_utility.hpp>
#include <boost/graph/iteration_macros.hpp>
#include <boost/graph/properties.hpp>
#include <boost/property_map/shared_array_property_map.hpp>

namespace boost
{

namespace detail
{

    // Traits associated with common subgraphs, used mainly to keep a
    // consistent type for the correspondence maps.
    template < typename GraphFirst, typename GraphSecond,
        typename VertexIndexMapFirst, typename VertexIndexMapSecond >
    struct mcgregor_common_subgraph_traits
    {
        typedef typename graph_traits< GraphFirst >::vertex_descriptor
            vertex_first_type;
        typedef typename graph_traits< GraphSecond >::vertex_descriptor
            vertex_second_type;

        typedef shared_array_property_map< vertex_second_type,
            VertexIndexMapFirst >
            correspondence_map_first_to_second_type;

        typedef shared_array_property_map< vertex_first_type,
            VertexIndexMapSecond >
            correspondence_map_second_to_first_type;
    };

} // namespace detail

// ==========================================================================

// Binary function object that returns true if the values for item1
// in property_map1 and item2 in property_map2 are equivalent.
template < typename PropertyMapFirst, typename PropertyMapSecond >
struct property_map_equivalent
{

    property_map_equivalent(const PropertyMapFirst property_map1,
        const PropertyMapSecond property_map2)
    : m_property_map1(property_map1), m_property_map2(property_map2)
    {
    }

    template < typename ItemFirst, typename ItemSecond >
    bool operator()(const ItemFirst item1, const ItemSecond item2)
    {
        return (get(m_property_map1, item1) == get(m_property_map2, item2));
    }

private:
    const PropertyMapFirst m_property_map1;
    const PropertyMapSecond m_property_map2;
};

// Returns a property_map_equivalent object that compares the values
// of property_map1 and property_map2.
template < typename PropertyMapFirst, typename PropertyMapSecond >
property_map_equivalent< PropertyMapFirst, PropertyMapSecond >
make_property_map_equivalent(
    const PropertyMapFirst property_map1, const PropertyMapSecond property_map2)
{

    return (property_map_equivalent< PropertyMapFirst, PropertyMapSecond >(
        property_map1, property_map2));
}

// Binary function object that always returns true.  Used when
// vertices or edges are always equivalent (i.e. have no labels).
struct always_equivalent
{

    template < typename ItemFirst, typename ItemSecond >
    bool operator()(const ItemFirst&, const ItemSecond&)
    {
        return (true);
    }
};

// ==========================================================================

namespace detail
{

    // Return true if new_vertex1 and new_vertex2 can extend the
    // subgraph represented by correspondence_map_1_to_2 and
    // correspondence_map_2_to_1.  The vertices_equivalent and
    // edges_equivalent predicates are used to test vertex and edge
    // equivalency between the two graphs.
    template < typename GraphFirst, typename GraphSecond,
        typename CorrespondenceMapFirstToSecond,
        typename CorrespondenceMapSecondToFirst,
        typename EdgeEquivalencePredicate, typename VertexEquivalencePredicate >
    bool can_extend_graph(const GraphFirst& graph1, const GraphSecond& graph2,
        CorrespondenceMapFirstToSecond correspondence_map_1_to_2,
        CorrespondenceMapSecondToFirst /*correspondence_map_2_to_1*/,
        typename graph_traits< GraphFirst >::vertices_size_type subgraph_size,
        typename graph_traits< GraphFirst >::vertex_descriptor new_vertex1,
        typename graph_traits< GraphSecond >::vertex_descriptor new_vertex2,
        EdgeEquivalencePredicate edges_equivalent,
        VertexEquivalencePredicate vertices_equivalent,
        bool only_connected_subgraphs)
    {
        typedef typename graph_traits< GraphSecond >::vertex_descriptor
            VertexSecond;

        typedef typename graph_traits< GraphFirst >::edge_descriptor EdgeFirst;
        typedef
            typename graph_traits< GraphSecond >::edge_descriptor EdgeSecond;

        // Check vertex equality
        if (!vertices_equivalent(new_vertex1, new_vertex2))
        {
            return (false);
        }

        // Vertices match and graph is empty, so we can extend the subgraph
        if (subgraph_size == 0)
        {
            return (true);
        }

        bool has_one_edge = false;

        // Verify edges with existing sub-graph
        BGL_FORALL_VERTICES_T(existing_vertex1, graph1, GraphFirst)
        {

            VertexSecond existing_vertex2
                = get(correspondence_map_1_to_2, existing_vertex1);

            // Skip unassociated vertices
            if (existing_vertex2 == graph_traits< GraphSecond >::null_vertex())
            {
                continue;
            }

            // NOTE: This will not work with parallel edges, since the
            // first matching edge is always chosen.
            EdgeFirst edge_to_new1, edge_from_new1;
            bool edge_to_new_exists1 = false, edge_from_new_exists1 = false;

            EdgeSecond edge_to_new2, edge_from_new2;
            bool edge_to_new_exists2 = false, edge_from_new_exists2 = false;

            // Search for edge from existing to new vertex (graph1)
            BGL_FORALL_OUTEDGES_T(existing_vertex1, edge1, graph1, GraphFirst)
            {
                if (target(edge1, graph1) == new_vertex1)
                {
                    edge_to_new1 = edge1;
                    edge_to_new_exists1 = true;
                    break;
                }
            }

            // Search for edge from existing to new vertex (graph2)
            BGL_FORALL_OUTEDGES_T(existing_vertex2, edge2, graph2, GraphSecond)
            {
                if (target(edge2, graph2) == new_vertex2)
                {
                    edge_to_new2 = edge2;
                    edge_to_new_exists2 = true;
                    break;
                }
            }

            // Make sure edges from existing to new vertices are equivalent
            if ((edge_to_new_exists1 != edge_to_new_exists2)
                || ((edge_to_new_exists1 && edge_to_new_exists2)
                    && !edges_equivalent(edge_to_new1, edge_to_new2)))
            {

                return (false);
            }

            bool is_undirected1 = is_undirected(graph1),
                 is_undirected2 = is_undirected(graph2);

            if (is_undirected1 && is_undirected2)
            {

                // Edge in both graphs exists and both graphs are undirected
                if (edge_to_new_exists1 && edge_to_new_exists2)
                {
                    has_one_edge = true;
                }

                continue;
            }
            else
            {

                if (!is_undirected1)
                {

                    // Search for edge from new to existing vertex (graph1)
                    BGL_FORALL_OUTEDGES_T(
                        new_vertex1, edge1, graph1, GraphFirst)
                    {
                        if (target(edge1, graph1) == existing_vertex1)
                        {
                            edge_from_new1 = edge1;
                            edge_from_new_exists1 = true;
                            break;
                        }
                    }
                }

                if (!is_undirected2)
                {

                    // Search for edge from new to existing vertex (graph2)
                    BGL_FORALL_OUTEDGES_T(
                        new_vertex2, edge2, graph2, GraphSecond)
                    {
                        if (target(edge2, graph2) == existing_vertex2)
                        {
                            edge_from_new2 = edge2;
                            edge_from_new_exists2 = true;
                            break;
                        }
                    }
                }

                // Make sure edges from new to existing vertices are equivalent
                if ((edge_from_new_exists1 != edge_from_new_exists2)
                    || ((edge_from_new_exists1 && edge_from_new_exists2)
                        && !edges_equivalent(edge_from_new1, edge_from_new2)))
                {

                    return (false);
                }

                if ((edge_from_new_exists1 && edge_from_new_exists2)
                    || (edge_to_new_exists1 && edge_to_new_exists2))
                {
                    has_one_edge = true;
                }

            } // else

        } // BGL_FORALL_VERTICES_T

        // Make sure new vertices are connected to the existing subgraph
        if (only_connected_subgraphs && !has_one_edge)
        {
            return (false);
        }

        return (true);
    }

    // Recursive method that does a depth-first search in the space of
    // potential subgraphs.  At each level, every new vertex pair from
    // both graphs is tested to see if it can extend the current
    // subgraph.  If so, the subgraph is output to subgraph_callback
    // in the form of two correspondence maps (one for each graph).
    // Returning false from subgraph_callback will terminate the
    // search.  Function returns true if the entire search space was
    // explored.
    template < typename GraphFirst, typename GraphSecond,
        typename VertexIndexMapFirst, typename VertexIndexMapSecond,
        typename CorrespondenceMapFirstToSecond,
        typename CorrespondenceMapSecondToFirst, typename VertexStackFirst,
        typename EdgeEquivalencePredicate, typename VertexEquivalencePredicate,
        typename SubGraphInternalCallback >
    bool mcgregor_common_subgraphs_internal(const GraphFirst& graph1,
        const GraphSecond& graph2, const VertexIndexMapFirst& vindex_map1,
        const VertexIndexMapSecond& vindex_map2,
        CorrespondenceMapFirstToSecond correspondence_map_1_to_2,
        CorrespondenceMapSecondToFirst correspondence_map_2_to_1,
        VertexStackFirst& vertex_stack1,
        EdgeEquivalencePredicate edges_equivalent,
        VertexEquivalencePredicate vertices_equivalent,
        bool only_connected_subgraphs,
        SubGraphInternalCallback subgraph_callback)
    {
        typedef
            typename graph_traits< GraphFirst >::vertex_descriptor VertexFirst;
        typedef typename graph_traits< GraphSecond >::vertex_descriptor
            VertexSecond;
        typedef typename graph_traits< GraphFirst >::vertices_size_type
            VertexSizeFirst;

        // Get iterators for vertices from both graphs
        typename graph_traits< GraphFirst >::vertex_iterator vertex1_iter,
            vertex1_end;

        typename graph_traits< GraphSecond >::vertex_iterator vertex2_begin,
            vertex2_end, vertex2_iter;

        boost::tie(vertex1_iter, vertex1_end) = vertices(graph1);
        boost::tie(vertex2_begin, vertex2_end) = vertices(graph2);
        vertex2_iter = vertex2_begin;

        // Iterate until all vertices have been visited
        BGL_FORALL_VERTICES_T(new_vertex1, graph1, GraphFirst)
        {

            VertexSecond existing_vertex2
                = get(correspondence_map_1_to_2, new_vertex1);

            // Skip already matched vertices in first graph
            if (existing_vertex2 != graph_traits< GraphSecond >::null_vertex())
            {
                continue;
            }

            BGL_FORALL_VERTICES_T(new_vertex2, graph2, GraphSecond)
            {

                VertexFirst existing_vertex1
                    = get(correspondence_map_2_to_1, new_vertex2);

                // Skip already matched vertices in second graph
                if (existing_vertex1
                    != graph_traits< GraphFirst >::null_vertex())
                {
                    continue;
                }

                // Check if current sub-graph can be extended with the matched
                // vertex pair
                if (can_extend_graph(graph1, graph2, correspondence_map_1_to_2,
                        correspondence_map_2_to_1,
                        (VertexSizeFirst)vertex_stack1.size(), new_vertex1,
                        new_vertex2, edges_equivalent, vertices_equivalent,
                        only_connected_subgraphs))
                {

                    // Keep track of old graph size for restoring later
                    VertexSizeFirst old_graph_size
                        = (VertexSizeFirst)vertex_stack1.size(),
                        new_graph_size = old_graph_size + 1;

                    // Extend subgraph
                    put(correspondence_map_1_to_2, new_vertex1, new_vertex2);
                    put(correspondence_map_2_to_1, new_vertex2, new_vertex1);
                    vertex_stack1.push(new_vertex1);

                    // Returning false from the callback will cancel iteration
                    if (!subgraph_callback(correspondence_map_1_to_2,
                            correspondence_map_2_to_1, new_graph_size))
                    {
                        return (false);
                    }

                    // Depth-first search into the state space of possible
                    // sub-graphs
                    bool continue_iteration
                        = mcgregor_common_subgraphs_internal(graph1, graph2,
                            vindex_map1, vindex_map2, correspondence_map_1_to_2,
                            correspondence_map_2_to_1, vertex_stack1,
                            edges_equivalent, vertices_equivalent,
                            only_connected_subgraphs, subgraph_callback);

                    if (!continue_iteration)
                    {
                        return (false);
                    }

                    // Restore previous state
                    if (vertex_stack1.size() > old_graph_size)
                    {

                        VertexFirst stack_vertex1 = vertex_stack1.top();
                        VertexSecond stack_vertex2
                            = get(correspondence_map_1_to_2, stack_vertex1);

                        // Contract subgraph
                        put(correspondence_map_1_to_2, stack_vertex1,
                            graph_traits< GraphSecond >::null_vertex());

                        put(correspondence_map_2_to_1, stack_vertex2,
                            graph_traits< GraphFirst >::null_vertex());

                        vertex_stack1.pop();
                    }

                } // if can_extend_graph

            } // BGL_FORALL_VERTICES_T (graph2)

        } // BGL_FORALL_VERTICES_T (graph1)

        return (true);
    }

    // Internal method that initializes blank correspondence maps and
    // a vertex stack for use in mcgregor_common_subgraphs_internal.
    template < typename GraphFirst, typename GraphSecond,
        typename VertexIndexMapFirst, typename VertexIndexMapSecond,
        typename EdgeEquivalencePredicate, typename VertexEquivalencePredicate,
        typename SubGraphInternalCallback >
    inline void mcgregor_common_subgraphs_internal_init(
        const GraphFirst& graph1, const GraphSecond& graph2,
        const VertexIndexMapFirst vindex_map1,
        const VertexIndexMapSecond vindex_map2,
        EdgeEquivalencePredicate edges_equivalent,
        VertexEquivalencePredicate vertices_equivalent,
        bool only_connected_subgraphs,
        SubGraphInternalCallback subgraph_callback)
    {
        typedef mcgregor_common_subgraph_traits< GraphFirst, GraphSecond,
            VertexIndexMapFirst, VertexIndexMapSecond >
            SubGraphTraits;

        typename SubGraphTraits::correspondence_map_first_to_second_type
            correspondence_map_1_to_2(num_vertices(graph1), vindex_map1);

        BGL_FORALL_VERTICES_T(vertex1, graph1, GraphFirst)
        {
            put(correspondence_map_1_to_2, vertex1,
                graph_traits< GraphSecond >::null_vertex());
        }

        typename SubGraphTraits::correspondence_map_second_to_first_type
            correspondence_map_2_to_1(num_vertices(graph2), vindex_map2);

        BGL_FORALL_VERTICES_T(vertex2, graph2, GraphSecond)
        {
            put(correspondence_map_2_to_1, vertex2,
                graph_traits< GraphFirst >::null_vertex());
        }

        typedef
            typename graph_traits< GraphFirst >::vertex_descriptor VertexFirst;

        std::stack< VertexFirst > vertex_stack1;

        mcgregor_common_subgraphs_internal(graph1, graph2, vindex_map1,
            vindex_map2, correspondence_map_1_to_2, correspondence_map_2_to_1,
            vertex_stack1, edges_equivalent, vertices_equivalent,
            only_connected_subgraphs, subgraph_callback);
    }

} // namespace detail

// ==========================================================================

// Enumerates all common subgraphs present in graph1 and graph2.
// Continues until the search space has been fully explored or false
// is returned from user_callback.
template < typename GraphFirst, typename GraphSecond,
    typename VertexIndexMapFirst, typename VertexIndexMapSecond,
    typename EdgeEquivalencePredicate, typename VertexEquivalencePredicate,
    typename SubGraphCallback >
void mcgregor_common_subgraphs(const GraphFirst& graph1,
    const GraphSecond& graph2, const VertexIndexMapFirst vindex_map1,
    const VertexIndexMapSecond vindex_map2,
    EdgeEquivalencePredicate edges_equivalent,
    VertexEquivalencePredicate vertices_equivalent,
    bool only_connected_subgraphs, SubGraphCallback user_callback)
{

    detail::mcgregor_common_subgraphs_internal_init(graph1, graph2, vindex_map1,
        vindex_map2, edges_equivalent, vertices_equivalent,
        only_connected_subgraphs, user_callback);
}

// Variant of mcgregor_common_subgraphs with all default parameters
template < typename GraphFirst, typename GraphSecond,
    typename SubGraphCallback >
void mcgregor_common_subgraphs(const GraphFirst& graph1,
    const GraphSecond& graph2, bool only_connected_subgraphs,
    SubGraphCallback user_callback)
{

    detail::mcgregor_common_subgraphs_internal_init(graph1, graph2,
        get(vertex_index, graph1), get(vertex_index, graph2),
        always_equivalent(), always_equivalent(), only_connected_subgraphs,
        user_callback);
}

// Named parameter variant of mcgregor_common_subgraphs
template < typename GraphFirst, typename GraphSecond, typename SubGraphCallback,
    typename Param, typename Tag, typename Rest >
void mcgregor_common_subgraphs(const GraphFirst& graph1,
    const GraphSecond& graph2, bool only_connected_subgraphs,
    SubGraphCallback user_callback,
    const bgl_named_params< Param, Tag, Rest >& params)
{

    detail::mcgregor_common_subgraphs_internal_init(graph1, graph2,
        choose_const_pmap(
            get_param(params, vertex_index1), graph1, vertex_index),
        choose_const_pmap(
            get_param(params, vertex_index2), graph2, vertex_index),
        choose_param(
            get_param(params, edges_equivalent_t()), always_equivalent()),
        choose_param(
            get_param(params, vertices_equivalent_t()), always_equivalent()),
        only_connected_subgraphs, user_callback);
}

// ==========================================================================

namespace detail
{

    // Binary function object that intercepts subgraphs from
    // mcgregor_common_subgraphs_internal and maintains a cache of
    // unique subgraphs.  The user callback is invoked for each unique
    // subgraph.
    template < typename GraphFirst, typename GraphSecond,
        typename VertexIndexMapFirst, typename VertexIndexMapSecond,
        typename SubGraphCallback >
    struct unique_subgraph_interceptor
    {

        typedef typename graph_traits< GraphFirst >::vertices_size_type
            VertexSizeFirst;

        typedef mcgregor_common_subgraph_traits< GraphFirst, GraphSecond,
            VertexIndexMapFirst, VertexIndexMapSecond >
            SubGraphTraits;

        typedef typename SubGraphTraits::correspondence_map_first_to_second_type
            CachedCorrespondenceMapFirstToSecond;

        typedef typename SubGraphTraits::correspondence_map_second_to_first_type
            CachedCorrespondenceMapSecondToFirst;

        typedef std::pair< VertexSizeFirst,
            std::pair< CachedCorrespondenceMapFirstToSecond,
                CachedCorrespondenceMapSecondToFirst > >
            SubGraph;

        typedef std::vector< SubGraph > SubGraphList;

        unique_subgraph_interceptor(const GraphFirst& graph1,
            const GraphSecond& graph2, const VertexIndexMapFirst vindex_map1,
            const VertexIndexMapSecond vindex_map2,
            SubGraphCallback user_callback)
        : m_graph1(graph1)
        , m_graph2(graph2)
        , m_vindex_map1(vindex_map1)
        , m_vindex_map2(vindex_map2)
        , m_subgraphs(make_shared< SubGraphList >())
        , m_user_callback(user_callback)
        {
        }

        template < typename CorrespondenceMapFirstToSecond,
            typename CorrespondenceMapSecondToFirst >
        bool operator()(
            CorrespondenceMapFirstToSecond correspondence_map_1_to_2,
            CorrespondenceMapSecondToFirst correspondence_map_2_to_1,
            VertexSizeFirst subgraph_size)
        {

            for (typename SubGraphList::const_iterator subgraph_iter
                 = m_subgraphs->begin();
                 subgraph_iter != m_subgraphs->end(); ++subgraph_iter)
            {

                SubGraph subgraph_cached = *subgraph_iter;

                // Compare subgraph sizes
                if (subgraph_size != subgraph_cached.first)
                {
                    continue;
                }

                if (!are_property_maps_different(correspondence_map_1_to_2,
                        subgraph_cached.second.first, m_graph1))
                {

                    // New subgraph is a duplicate
                    return (true);
                }
            }

            // Subgraph is unique, so make a cached copy
            CachedCorrespondenceMapFirstToSecond new_subgraph_1_to_2
                = CachedCorrespondenceMapFirstToSecond(
                    num_vertices(m_graph1), m_vindex_map1);

            CachedCorrespondenceMapSecondToFirst new_subgraph_2_to_1
                = CorrespondenceMapSecondToFirst(
                    num_vertices(m_graph2), m_vindex_map2);

            BGL_FORALL_VERTICES_T(vertex1, m_graph1, GraphFirst)
            {
                put(new_subgraph_1_to_2, vertex1,
                    get(correspondence_map_1_to_2, vertex1));
            }

            BGL_FORALL_VERTICES_T(vertex2, m_graph2, GraphFirst)
            {
                put(new_subgraph_2_to_1, vertex2,
                    get(correspondence_map_2_to_1, vertex2));
            }

            m_subgraphs->push_back(std::make_pair(subgraph_size,
                std::make_pair(new_subgraph_1_to_2, new_subgraph_2_to_1)));

            return (m_user_callback(correspondence_map_1_to_2,
                correspondence_map_2_to_1, subgraph_size));
        }

    private:
        const GraphFirst& m_graph1;
        const GraphFirst& m_graph2;
        const VertexIndexMapFirst m_vindex_map1;
        const VertexIndexMapSecond m_vindex_map2;
        shared_ptr< SubGraphList > m_subgraphs;
        SubGraphCallback m_user_callback;
    };

} // namespace detail

// Enumerates all unique common subgraphs between graph1 and graph2.
// The user callback is invoked for each unique subgraph as they are
// discovered.
template < typename GraphFirst, typename GraphSecond,
    typename VertexIndexMapFirst, typename VertexIndexMapSecond,
    typename EdgeEquivalencePredicate, typename VertexEquivalencePredicate,
    typename SubGraphCallback >
void mcgregor_common_subgraphs_unique(const GraphFirst& graph1,
    const GraphSecond& graph2, const VertexIndexMapFirst vindex_map1,
    const VertexIndexMapSecond vindex_map2,
    EdgeEquivalencePredicate edges_equivalent,
    VertexEquivalencePredicate vertices_equivalent,
    bool only_connected_subgraphs, SubGraphCallback user_callback)
{
    detail::unique_subgraph_interceptor< GraphFirst, GraphSecond,
        VertexIndexMapFirst, VertexIndexMapSecond, SubGraphCallback >
        unique_callback(
            graph1, graph2, vindex_map1, vindex_map2, user_callback);

    detail::mcgregor_common_subgraphs_internal_init(graph1, graph2, vindex_map1,
        vindex_map2, edges_equivalent, vertices_equivalent,
        only_connected_subgraphs, unique_callback);
}

// Variant of mcgregor_common_subgraphs_unique with all default
// parameters.
template < typename GraphFirst, typename GraphSecond,
    typename SubGraphCallback >
void mcgregor_common_subgraphs_unique(const GraphFirst& graph1,
    const GraphSecond& graph2, bool only_connected_subgraphs,
    SubGraphCallback user_callback)
{
    mcgregor_common_subgraphs_unique(graph1, graph2, get(vertex_index, graph1),
        get(vertex_index, graph2), always_equivalent(), always_equivalent(),
        only_connected_subgraphs, user_callback);
}

// Named parameter variant of mcgregor_common_subgraphs_unique
template < typename GraphFirst, typename GraphSecond, typename SubGraphCallback,
    typename Param, typename Tag, typename Rest >
void mcgregor_common_subgraphs_unique(const GraphFirst& graph1,
    const GraphSecond& graph2, bool only_connected_subgraphs,
    SubGraphCallback user_callback,
    const bgl_named_params< Param, Tag, Rest >& params)
{
    mcgregor_common_subgraphs_unique(graph1, graph2,
        choose_const_pmap(
            get_param(params, vertex_index1), graph1, vertex_index),
        choose_const_pmap(
            get_param(params, vertex_index2), graph2, vertex_index),
        choose_param(
            get_param(params, edges_equivalent_t()), always_equivalent()),
        choose_param(
            get_param(params, vertices_equivalent_t()), always_equivalent()),
        only_connected_subgraphs, user_callback);
}

// ==========================================================================

namespace detail
{

    // Binary function object that intercepts subgraphs from
    // mcgregor_common_subgraphs_internal and maintains a cache of the
    // largest subgraphs.
    template < typename GraphFirst, typename GraphSecond,
        typename VertexIndexMapFirst, typename VertexIndexMapSecond,
        typename SubGraphCallback >
    struct maximum_subgraph_interceptor
    {

        typedef typename graph_traits< GraphFirst >::vertices_size_type
            VertexSizeFirst;

        typedef mcgregor_common_subgraph_traits< GraphFirst, GraphSecond,
            VertexIndexMapFirst, VertexIndexMapSecond >
            SubGraphTraits;

        typedef typename SubGraphTraits::correspondence_map_first_to_second_type
            CachedCorrespondenceMapFirstToSecond;

        typedef typename SubGraphTraits::correspondence_map_second_to_first_type
            CachedCorrespondenceMapSecondToFirst;

        typedef std::pair< VertexSizeFirst,
            std::pair< CachedCorrespondenceMapFirstToSecond,
                CachedCorrespondenceMapSecondToFirst > >
            SubGraph;

        typedef std::vector< SubGraph > SubGraphList;

        maximum_subgraph_interceptor(const GraphFirst& graph1,
            const GraphSecond& graph2, const VertexIndexMapFirst vindex_map1,
            const VertexIndexMapSecond vindex_map2,
            SubGraphCallback user_callback)
        : m_graph1(graph1)
        , m_graph2(graph2)
        , m_vindex_map1(vindex_map1)
        , m_vindex_map2(vindex_map2)
        , m_subgraphs(make_shared< SubGraphList >())
        , m_largest_size_so_far(make_shared< VertexSizeFirst >(0))
        , m_user_callback(user_callback)
        {
        }

        template < typename CorrespondenceMapFirstToSecond,
            typename CorrespondenceMapSecondToFirst >
        bool operator()(
            CorrespondenceMapFirstToSecond correspondence_map_1_to_2,
            CorrespondenceMapSecondToFirst correspondence_map_2_to_1,
            VertexSizeFirst subgraph_size)
        {

            if (subgraph_size > *m_largest_size_so_far)
            {
                m_subgraphs->clear();
                *m_largest_size_so_far = subgraph_size;
            }

            if (subgraph_size == *m_largest_size_so_far)
            {

                // Make a cached copy
                CachedCorrespondenceMapFirstToSecond new_subgraph_1_to_2
                    = CachedCorrespondenceMapFirstToSecond(
                        num_vertices(m_graph1), m_vindex_map1);

                CachedCorrespondenceMapSecondToFirst new_subgraph_2_to_1
                    = CachedCorrespondenceMapSecondToFirst(
                        num_vertices(m_graph2), m_vindex_map2);

                BGL_FORALL_VERTICES_T(vertex1, m_graph1, GraphFirst)
                {
                    put(new_subgraph_1_to_2, vertex1,
                        get(correspondence_map_1_to_2, vertex1));
                }

                BGL_FORALL_VERTICES_T(vertex2, m_graph2, GraphFirst)
                {
                    put(new_subgraph_2_to_1, vertex2,
                        get(correspondence_map_2_to_1, vertex2));
                }

                m_subgraphs->push_back(std::make_pair(subgraph_size,
                    std::make_pair(new_subgraph_1_to_2, new_subgraph_2_to_1)));
            }

            return (true);
        }

        void output_subgraphs()
        {
            for (typename SubGraphList::const_iterator subgraph_iter
                 = m_subgraphs->begin();
                 subgraph_iter != m_subgraphs->end(); ++subgraph_iter)
            {

                SubGraph subgraph_cached = *subgraph_iter;
                m_user_callback(subgraph_cached.second.first,
                    subgraph_cached.second.second, subgraph_cached.first);
            }
        }

    private:
        const GraphFirst& m_graph1;
        const GraphFirst& m_graph2;
        const VertexIndexMapFirst m_vindex_map1;
        const VertexIndexMapSecond m_vindex_map2;
        shared_ptr< SubGraphList > m_subgraphs;
        shared_ptr< VertexSizeFirst > m_largest_size_so_far;
        SubGraphCallback m_user_callback;
    };

} // namespace detail

// Enumerates the largest common subgraphs found between graph1
// and graph2.  Note that the ENTIRE search space is explored before
// user_callback is actually invoked.
template < typename GraphFirst, typename GraphSecond,
    typename VertexIndexMapFirst, typename VertexIndexMapSecond,
    typename EdgeEquivalencePredicate, typename VertexEquivalencePredicate,
    typename SubGraphCallback >
void mcgregor_common_subgraphs_maximum(const GraphFirst& graph1,
    const GraphSecond& graph2, const VertexIndexMapFirst vindex_map1,
    const VertexIndexMapSecond vindex_map2,
    EdgeEquivalencePredicate edges_equivalent,
    VertexEquivalencePredicate vertices_equivalent,
    bool only_connected_subgraphs, SubGraphCallback user_callback)
{
    detail::maximum_subgraph_interceptor< GraphFirst, GraphSecond,
        VertexIndexMapFirst, VertexIndexMapSecond, SubGraphCallback >
        max_interceptor(
            graph1, graph2, vindex_map1, vindex_map2, user_callback);

    detail::mcgregor_common_subgraphs_internal_init(graph1, graph2, vindex_map1,
        vindex_map2, edges_equivalent, vertices_equivalent,
        only_connected_subgraphs, max_interceptor);

    // Only output the largest subgraphs
    max_interceptor.output_subgraphs();
}

// Variant of mcgregor_common_subgraphs_maximum with all default
// parameters.
template < typename GraphFirst, typename GraphSecond,
    typename SubGraphCallback >
void mcgregor_common_subgraphs_maximum(const GraphFirst& graph1,
    const GraphSecond& graph2, bool only_connected_subgraphs,
    SubGraphCallback user_callback)
{
    mcgregor_common_subgraphs_maximum(graph1, graph2, get(vertex_index, graph1),
        get(vertex_index, graph2), always_equivalent(), always_equivalent(),
        only_connected_subgraphs, user_callback);
}

// Named parameter variant of mcgregor_common_subgraphs_maximum
template < typename GraphFirst, typename GraphSecond, typename SubGraphCallback,
    typename Param, typename Tag, typename Rest >
void mcgregor_common_subgraphs_maximum(const GraphFirst& graph1,
    const GraphSecond& graph2, bool only_connected_subgraphs,
    SubGraphCallback user_callback,
    const bgl_named_params< Param, Tag, Rest >& params)
{
    mcgregor_common_subgraphs_maximum(graph1, graph2,
        choose_const_pmap(
            get_param(params, vertex_index1), graph1, vertex_index),
        choose_const_pmap(
            get_param(params, vertex_index2), graph2, vertex_index),
        choose_param(
            get_param(params, edges_equivalent_t()), always_equivalent()),
        choose_param(
            get_param(params, vertices_equivalent_t()), always_equivalent()),
        only_connected_subgraphs, user_callback);
}

// ==========================================================================

namespace detail
{

    // Binary function object that intercepts subgraphs from
    // mcgregor_common_subgraphs_internal and maintains a cache of the
    // largest, unique subgraphs.
    template < typename GraphFirst, typename GraphSecond,
        typename VertexIndexMapFirst, typename VertexIndexMapSecond,
        typename SubGraphCallback >
    struct unique_maximum_subgraph_interceptor
    {

        typedef typename graph_traits< GraphFirst >::vertices_size_type
            VertexSizeFirst;

        typedef mcgregor_common_subgraph_traits< GraphFirst, GraphSecond,
            VertexIndexMapFirst, VertexIndexMapSecond >
            SubGraphTraits;

        typedef typename SubGraphTraits::correspondence_map_first_to_second_type
            CachedCorrespondenceMapFirstToSecond;

        typedef typename SubGraphTraits::correspondence_map_second_to_first_type
            CachedCorrespondenceMapSecondToFirst;

        typedef std::pair< VertexSizeFirst,
            std::pair< CachedCorrespondenceMapFirstToSecond,
                CachedCorrespondenceMapSecondToFirst > >
            SubGraph;

        typedef std::vector< SubGraph > SubGraphList;

        unique_maximum_subgraph_interceptor(const GraphFirst& graph1,
            const GraphSecond& graph2, const VertexIndexMapFirst vindex_map1,
            const VertexIndexMapSecond vindex_map2,
            SubGraphCallback user_callback)
        : m_graph1(graph1)
        , m_graph2(graph2)
        , m_vindex_map1(vindex_map1)
        , m_vindex_map2(vindex_map2)
        , m_subgraphs(make_shared< SubGraphList >())
        , m_largest_size_so_far(make_shared< VertexSizeFirst >(0))
        , m_user_callback(user_callback)
        {
        }

        template < typename CorrespondenceMapFirstToSecond,
            typename CorrespondenceMapSecondToFirst >
        bool operator()(
            CorrespondenceMapFirstToSecond correspondence_map_1_to_2,
            CorrespondenceMapSecondToFirst correspondence_map_2_to_1,
            VertexSizeFirst subgraph_size)
        {

            if (subgraph_size > *m_largest_size_so_far)
            {
                m_subgraphs->clear();
                *m_largest_size_so_far = subgraph_size;
            }

            if (subgraph_size == *m_largest_size_so_far)
            {

                // Check if subgraph is unique
                for (typename SubGraphList::const_iterator subgraph_iter
                     = m_subgraphs->begin();
                     subgraph_iter != m_subgraphs->end(); ++subgraph_iter)
                {

                    SubGraph subgraph_cached = *subgraph_iter;

                    if (!are_property_maps_different(correspondence_map_1_to_2,
                            subgraph_cached.second.first, m_graph1))
                    {

                        // New subgraph is a duplicate
                        return (true);
                    }
                }

                // Subgraph is unique, so make a cached copy
                CachedCorrespondenceMapFirstToSecond new_subgraph_1_to_2
                    = CachedCorrespondenceMapFirstToSecond(
                        num_vertices(m_graph1), m_vindex_map1);

                CachedCorrespondenceMapSecondToFirst new_subgraph_2_to_1
                    = CachedCorrespondenceMapSecondToFirst(
                        num_vertices(m_graph2), m_vindex_map2);

                BGL_FORALL_VERTICES_T(vertex1, m_graph1, GraphFirst)
                {
                    put(new_subgraph_1_to_2, vertex1,
                        get(correspondence_map_1_to_2, vertex1));
                }

                BGL_FORALL_VERTICES_T(vertex2, m_graph2, GraphFirst)
                {
                    put(new_subgraph_2_to_1, vertex2,
                        get(correspondence_map_2_to_1, vertex2));
                }

                m_subgraphs->push_back(std::make_pair(subgraph_size,
                    std::make_pair(new_subgraph_1_to_2, new_subgraph_2_to_1)));
            }

            return (true);
        }

        void output_subgraphs()
        {
            for (typename SubGraphList::const_iterator subgraph_iter
                 = m_subgraphs->begin();
                 subgraph_iter != m_subgraphs->end(); ++subgraph_iter)
            {

                SubGraph subgraph_cached = *subgraph_iter;
                m_user_callback(subgraph_cached.second.first,
                    subgraph_cached.second.second, subgraph_cached.first);
            }
        }

    private:
        const GraphFirst& m_graph1;
        const GraphFirst& m_graph2;
        const VertexIndexMapFirst m_vindex_map1;
        const VertexIndexMapSecond m_vindex_map2;
        shared_ptr< SubGraphList > m_subgraphs;
        shared_ptr< VertexSizeFirst > m_largest_size_so_far;
        SubGraphCallback m_user_callback;
    };

} // namespace detail

// Enumerates the largest, unique common subgraphs found between
// graph1 and graph2.  Note that the ENTIRE search space is explored
// before user_callback is actually invoked.
template < typename GraphFirst, typename GraphSecond,
    typename VertexIndexMapFirst, typename VertexIndexMapSecond,
    typename EdgeEquivalencePredicate, typename VertexEquivalencePredicate,
    typename SubGraphCallback >
void mcgregor_common_subgraphs_maximum_unique(const GraphFirst& graph1,
    const GraphSecond& graph2, const VertexIndexMapFirst vindex_map1,
    const VertexIndexMapSecond vindex_map2,
    EdgeEquivalencePredicate edges_equivalent,
    VertexEquivalencePredicate vertices_equivalent,
    bool only_connected_subgraphs, SubGraphCallback user_callback)
{
    detail::unique_maximum_subgraph_interceptor< GraphFirst, GraphSecond,
        VertexIndexMapFirst, VertexIndexMapSecond, SubGraphCallback >
        unique_max_interceptor(
            graph1, graph2, vindex_map1, vindex_map2, user_callback);

    detail::mcgregor_common_subgraphs_internal_init(graph1, graph2, vindex_map1,
        vindex_map2, edges_equivalent, vertices_equivalent,
        only_connected_subgraphs, unique_max_interceptor);

    // Only output the largest, unique subgraphs
    unique_max_interceptor.output_subgraphs();
}

// Variant of mcgregor_common_subgraphs_maximum_unique with all default
// parameters
template < typename GraphFirst, typename GraphSecond,
    typename SubGraphCallback >
void mcgregor_common_subgraphs_maximum_unique(const GraphFirst& graph1,
    const GraphSecond& graph2, bool only_connected_subgraphs,
    SubGraphCallback user_callback)
{

    mcgregor_common_subgraphs_maximum_unique(graph1, graph2,
        get(vertex_index, graph1), get(vertex_index, graph2),
        always_equivalent(), always_equivalent(), only_connected_subgraphs,
        user_callback);
}

// Named parameter variant of
// mcgregor_common_subgraphs_maximum_unique
template < typename GraphFirst, typename GraphSecond, typename SubGraphCallback,
    typename Param, typename Tag, typename Rest >
void mcgregor_common_subgraphs_maximum_unique(const GraphFirst& graph1,
    const GraphSecond& graph2, bool only_connected_subgraphs,
    SubGraphCallback user_callback,
    const bgl_named_params< Param, Tag, Rest >& params)
{
    mcgregor_common_subgraphs_maximum_unique(graph1, graph2,
        choose_const_pmap(
            get_param(params, vertex_index1), graph1, vertex_index),
        choose_const_pmap(
            get_param(params, vertex_index2), graph2, vertex_index),
        choose_param(
            get_param(params, edges_equivalent_t()), always_equivalent()),
        choose_param(
            get_param(params, vertices_equivalent_t()), always_equivalent()),
        only_connected_subgraphs, user_callback);
}

// ==========================================================================

// Fills a membership map (vertex -> bool) using the information
// present in correspondence_map_1_to_2. Every vertex in a
// membership map will have a true value only if it is not
// associated with a null vertex in the correspondence map.
template < typename GraphSecond, typename GraphFirst,
    typename CorrespondenceMapFirstToSecond, typename MembershipMapFirst >
void fill_membership_map(const GraphFirst& graph1,
    const CorrespondenceMapFirstToSecond correspondence_map_1_to_2,
    MembershipMapFirst membership_map1)
{

    BGL_FORALL_VERTICES_T(vertex1, graph1, GraphFirst)
    {
        put(membership_map1, vertex1,
            get(correspondence_map_1_to_2, vertex1)
                != graph_traits< GraphSecond >::null_vertex());
    }
}

// Traits associated with a membership map filtered graph.  Provided
// for convenience to access graph and vertex filter types.
template < typename Graph, typename MembershipMap >
struct membership_filtered_graph_traits
{
    typedef property_map_filter< MembershipMap > vertex_filter_type;
    typedef filtered_graph< Graph, keep_all, vertex_filter_type > graph_type;
};

// Returns a filtered sub-graph of graph whose edge and vertex
// inclusion is dictated by membership_map.
template < typename Graph, typename MembershipMap >
typename membership_filtered_graph_traits< Graph, MembershipMap >::graph_type
make_membership_filtered_graph(
    const Graph& graph, MembershipMap& membership_map)
{

    typedef membership_filtered_graph_traits< Graph, MembershipMap > MFGTraits;
    typedef typename MFGTraits::graph_type MembershipFilteredGraph;

    typename MFGTraits::vertex_filter_type v_filter(membership_map);

    return (MembershipFilteredGraph(graph, keep_all(), v_filter));
}

} // namespace boost

#endif // BOOST_GRAPH_MCGREGOR_COMMON_SUBGRAPHS_HPP


// __BEGIN_LICENSE__
// Copyright (C) 2006-2011 United States Government as represented by
// the Administrator of the National Aeronautics and Space Administration.
// All Rights Reserved.
// __END_LICENSE__

#include <vw/BundleAdjustment/ControlNetworkLoader.h>
#include <vw/Stereo/StereoModel.h>

using namespace vw;
using namespace vw::ba;

#include <boost/filesystem/fstream.hpp>

namespace fs = boost::filesystem;

struct ContainsEqualIP {
  ip::InterestPoint& m_compare;
  ContainsEqualIP( ip::InterestPoint& comp ) : m_compare(comp) {}

  bool operator()( boost::shared_ptr<IPFeature> in ) {
    if (  m_compare.x == in->m_ip.x &&
          m_compare.y == in->m_ip.y )
      return true;
    return false;
  }
};

// Utility for checking that the point is BA safe
void safe_measurement( ip::InterestPoint& ip ) {
  if ( ip.scale <= 0 ) ip.scale = 10;
}

void vw::ba::triangulate_control_point( ControlPoint& cp,
                                        std::vector<boost::shared_ptr<camera::CameraModel> > const& camera_models,
                                        double const& minimum_angle ) {
  Vector3 position_sum;
  double error = 0, error_sum = 0;
  size_t count = 0;

  // 4.1.) Building a listing of triangulation
  for ( size_t j = 0, k = 1; k < cp.size(); j++, k++ ) {
    // Make sure camera centers are not equal
    size_t j_cam_id = cp[j].image_id();
    size_t k_cam_id = cp[k].image_id();
    if ( norm_2( camera_models[j_cam_id]->camera_center( cp[j].position() ) -
                 camera_models[k_cam_id]->camera_center( cp[k].position() ) ) > 1e-6 ) {

      try {
        stereo::StereoModel sm( camera_models[ j_cam_id ].get(),
                                camera_models[ k_cam_id ].get() );

        if ( sm.convergence_angle( cp[j].position(),
                                   cp[k].position() ) >
             minimum_angle ) {
	  count++;
	  position_sum += sm( cp[j].position(), cp[k].position(), error );
          error_sum += error;
        }
      } catch ( const camera::PixelToRayErr& ) { /* Just let it go */ }
    }
  }

  // 4.2.) Summing, Averaging, and Storing
  if ( !count ) {
    vw_out(WarningMessage,"ba") << "Unable to triangulation position for point!\n";
    // At the very least we can provide a point that is some
    // distance out from the camera center and is in the 'general'
    // area.
    size_t j = cp[0].image_id();
    try {
      cp.set_position( camera_models[j]->camera_center(cp[0].position()) +
                       camera_models[j]->pixel_to_vector(cp[0].position())*10 );
    } catch ( const camera::PixelToRayErr& ) {
      cp.set_position( camera_models[j]->camera_center(cp[0].position()) +
                       camera_models[j]->camera_pose(cp[0].position()).rotate(Vector3(0,0,10)) );
    }
  } else {
    error_sum /= double(count);
    cp.set_position( position_sum / double(count) );
  }
}

void vw::ba::build_control_network( ba::ControlNetwork& cnet,
                                     std::vector<boost::shared_ptr<camera::CameraModel> > const& camera_models,
                                     std::vector<std::string> const& image_files,
                                     size_t min_matches,
                                     std::vector<std::string> const& directories ) {
  cnet.clear();

  // We can't guarantee that image_files is sorted, so we make a
  // std::map to give ourselves a sorted list and access to a binary
  // search.
  std::map<std::string,size_t> image_prefix_map;
  size_t count = 0;
  ba::CameraRelationNetwork<ba::IPFeature> crn;
  BOOST_FOREACH( std::string const& file, image_files ) {
    fs::path file_path(file);
    image_prefix_map[file_path.replace_extension().string()] = count;
    crn.add_node( ba::CameraNode<ba::IPFeature>( count,
                                                 file_path.stem() ) );
    count++;
  }

  // Searching through the directories available to us.
  size_t num_load_rejected = 0, num_loaded = 0;
  BOOST_FOREACH( std::string const& directory, directories ) {
    vw_out(VerboseDebugMessage,"ba") << "\tOpening directory \""
                                     << directory << "\".\n";
    fs::directory_iterator end_itr;
    for ( fs::directory_iterator obj( directory ); obj != end_itr; ++obj ) {

      // Skip if not a file
      if ( !fs::is_regular_file( obj->status() ) ) continue;
      // Skip if not a match file
      if ( obj->path().extension() != ".match" ) continue;

      // Pull out the prefixes that made up that match file
      std::string match_base = obj->path().stem();
      size_t split_pt = match_base.find("__");
      if ( split_pt == std::string::npos ) continue;
      std::string prefix1 = match_base.substr(0,split_pt);
      std::string prefix2 = match_base.substr(split_pt+2,match_base.size()-split_pt-2);

      // Extract the image indices that correspond to image prefixes.
      typedef std::map<std::string,size_t>::iterator MapIterator;
      MapIterator it1 = image_prefix_map.find( prefix1 );
      MapIterator it2 = image_prefix_map.find( prefix2 );
      if ( it1 == image_prefix_map.end() ||
           it2 == image_prefix_map.end() ) continue;

      // Actually read in the file as it seems we've found something correct
      std::vector<ip::InterestPoint> ip1, ip2;
      ip::read_binary_match_file( obj->string(), ip1, ip2 );
      if ( ip1.size() < min_matches ) {
        vw_out(VerboseDebugMessage,"ba") << "\t" << obj->string() << "    "
                                         << it1->second << " <-> " << it2->second << " : "
                                         << ip1.size() << " matches. [rejected]\n";
        num_load_rejected += ip1.size();
      } else {
        vw_out(VerboseDebugMessage,"ba") << "\t" << obj->string() << "    "
                                         << it1->second << " <-> " << it2->second << " : "
                                         << ip1.size() << " matches.\n";
        num_loaded += ip1.size();

        // Remove descriptors from interest points and correct scale
        std::for_each( ip1.begin(), ip1.end(), ip::remove_descriptor );
        std::for_each( ip2.begin(), ip2.end(), ip::remove_descriptor );
        std::for_each( ip1.begin(), ip1.end(), safe_measurement );
        std::for_each( ip2.begin(), ip2.end(), safe_measurement );

        typedef boost::shared_ptr< ba::IPFeature > f_ptr;
        typedef std::list< f_ptr >::iterator f_itr;

        // Checking to see if features already exist, adding if they
        // don't, then linking them.
        for ( size_t k = 0; k < ip1.size(); k++ ) {
          f_itr ipfeature1 = std::find_if( crn[it1->second].begin(),
                                           crn[it1->second].end(),
                                           ContainsEqualIP( ip1[k] ) );
          f_itr ipfeature2 = std::find_if( crn[it2->second].begin(),
                                           crn[it2->second].end(),
                                           ContainsEqualIP( ip2[k] ) );
          if ( ipfeature1 == crn[it1->second].end() ) {
            crn[it1->second].relations.push_front( f_ptr( new ba::IPFeature( ip1[k], it1->second ) ) );
            ipfeature1 = crn[it1->second].begin();
          }
          if ( ipfeature2 == crn[it2->second].end() ) {
            crn[it2->second].relations.push_front( f_ptr( new ba::IPFeature( ip2[k], it2->second ) ) );
            ipfeature2 = crn[it2->second].begin();
          }

          // Doubly linking
          (*ipfeature1)->connection( *ipfeature2, false );
          (*ipfeature2)->connection( *ipfeature1, false );
        }
      }
    }
  } // end search through directories
  if ( num_load_rejected != 0 ) {
    vw_out(WarningMessage,"ba") << "\tDidn't load " << num_load_rejected
                                << " matches due to inadequacy.\n";
    vw_out(WarningMessage,"ba") << "\tLoaded " << num_loaded << " matches.\n";
  }

  // Building control network
  crn.write_controlnetwork( cnet );

  // Triangulating Positions
  {
    TerminalProgressCallback progress("ba", "Triangulating:");
    progress.report_progress(0);
    double inc_prog = 1.0/double(cnet.size());
    double min_angle = 5.0*M_PI/180.0;
    BOOST_FOREACH( ba::ControlPoint& cpoint, cnet ) {
      progress.report_incremental_progress(inc_prog );
      ba::triangulate_control_point( cpoint, camera_models, min_angle );
    }
    progress.report_finished();
  }
}


/**
 *  This file is part of dvo.
 *
 *  Copyright 2012 Christian Kerl <christian.kerl@in.tum.de> (Technical University of Munich)
 *  For more information see <http://vision.in.tum.de/data/software/dvo>.
 *
 *  dvo 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.
 *
 *  dvo 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 dvo.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <iomanip>

#include <dvo/dense_tracking.h>
#include <dvo/dense_tracking_impl.h>

#include <assert.h>
#include <sophus/se3.hpp>

#include <Eigen/Core>

#include <dvo/core/datatypes.h>
#include <dvo/core/point_selection_predicates.h>
#include <dvo/util/revertable.h>
#include <dvo/util/stopwatch.h>
#include <dvo/util/id_generator.h>
#include <dvo/util/histogram.h>
//#include <dvo/visualization/visualizer.h>

namespace dvo
{

using namespace dvo::core;
using namespace dvo::util;

const DenseTracker::Config& DenseTracker::getDefaultConfig()
{
  static Config defaultConfig;

  return defaultConfig;
}

static const Eigen::IOFormat YamlArrayFmt(Eigen::FullPrecision, Eigen::DontAlignCols, ",", ",", "", "", "[", "]");

DenseTracker::DenseTracker(const Config& config) :
    itctx_(cfg),
    weight_calculation_(),
    selection_predicate_(),
    reference_selection_(selection_predicate_)
{
  configure(config);
}

DenseTracker::DenseTracker(const DenseTracker& other) :
  itctx_(cfg),
  weight_calculation_(),
  selection_predicate_(),
  reference_selection_(selection_predicate_)
{
  configure(other.configuration());
}

void DenseTracker::configure(const Config& config)
{
  assert(config.IsSane());

  cfg = config;

  selection_predicate_.intensity_threshold = cfg.IntensityDerivativeThreshold;
  selection_predicate_.depth_threshold = cfg.DepthDerivativeThreshold;

  if(cfg.UseWeighting)
  {
    weight_calculation_
      .scaleEstimator(ScaleEstimators::get(cfg.ScaleEstimatorType))
      .scaleEstimator()->configure(cfg.ScaleEstimatorParam);

    weight_calculation_
      .influenceFunction(InfluenceFunctions::get(cfg.InfluenceFuntionType))
      .influenceFunction()->configure(cfg.InfluenceFunctionParam);
  }
  else
  {
    weight_calculation_
      .scaleEstimator(ScaleEstimators::get(ScaleEstimators::Unit))
      .influenceFunction(InfluenceFunctions::get(InfluenceFunctions::Unit));
  }
}

bool DenseTracker::match(RgbdImagePyramid& reference, RgbdImagePyramid& current, Eigen::Affine3d& transformation)
{
  Result result;
  result.Transformation = transformation;

  bool success = match(reference, current, result);

  transformation = result.Transformation;

  return success;
}

bool DenseTracker::match(dvo::core::PointSelection& reference, RgbdImagePyramid& current, Eigen::Affine3d& transformation)
{
  Result result;
  result.Transformation = transformation;

  bool success = match(reference, current, result);

  transformation = result.Transformation;

  return success;
}

bool DenseTracker::match(dvo::core::RgbdImagePyramid& reference, dvo::core::RgbdImagePyramid& current, dvo::DenseTracker::Result& result)
{
  reference.compute(cfg.getNumLevels());
  reference_selection_.setRgbdImagePyramid(reference);

  return match(reference_selection_, current, result);
}

bool DenseTracker::match(dvo::core::PointSelection& reference, dvo::core::RgbdImagePyramid& current, dvo::DenseTracker::Result& result)
{
  current.compute(cfg.getNumLevels());

  bool success = true;

  if(cfg.UseInitialEstimate)
  {
    assert(!result.isNaN() && "Provided initialization is NaN!");
  }
  else
  {
    result.setIdentity();
  }

  // our first increment is the given guess
  Sophus::SE3d inc(result.Transformation.rotation(), result.Transformation.translation());

  Revertable<Sophus::SE3d> initial(inc);
  Revertable<Sophus::SE3d> estimate;

  bool accept = true;

  //static stopwatch_collection sw_level(5, "l", 100);
  //static stopwatch_collection sw_it(5, "it@l", 500);
  //static stopwatch_collection sw_error(5, "err@l", 500);
  //static stopwatch_collection sw_linsys(5, "linsys@l", 500);
  //static stopwatch_collection sw_prep(5, "prep@l", 100);

  if(points_error.size() < reference.getMaximumNumberOfPoints(cfg.LastLevel))
    points_error.resize(reference.getMaximumNumberOfPoints(cfg.LastLevel));
  if(residuals.size() < reference.getMaximumNumberOfPoints(cfg.LastLevel))
    residuals.resize(reference.getMaximumNumberOfPoints(cfg.LastLevel));
  if(weights.size() < reference.getMaximumNumberOfPoints(cfg.LastLevel))
    weights.resize(reference.getMaximumNumberOfPoints(cfg.LastLevel));

  std::vector<uint8_t> valid_residuals;

  bool debug = false;
  if(debug)
  {
    reference.debug(true);
    valid_residuals.resize(reference.getMaximumNumberOfPoints(cfg.LastLevel));
  }
  /*
  std::stringstream name;
  name << std::setiosflags(std::ios::fixed) << std::setprecision(2) << current.timestamp() << "_error.avi";

  cv::Size s = reference.getRgbdImagePyramid().level(size_t(cfg.LastLevel)).intensity.size();
  cv::Mat video_frame(s.height, s.width * 2, CV_32FC1), video_frame_u8;
  cv::VideoWriter vw(name.str(), CV_FOURCC('P','I','M','1'), 30, video_frame.size(), false);
  float rgb_max = 0.0;
  float depth_max = 0.0;

  std::stringstream name1;
  name1 << std::setiosflags(std::ios::fixed) << std::setprecision(2) << current.timestamp() << "_ref.png";

  cv::imwrite(name1.str(), current.level(0).rgb);

  std::stringstream name2;
  name2 << std::setiosflags(std::ios::fixed) << std::setprecision(2) << current.timestamp() << "_cur.png";

  cv::imwrite(name2.str(), reference.getRgbdImagePyramid().level(0).rgb);
  */
  Eigen::Vector2f mean;
  mean.setZero();
  Eigen::Matrix2f /*first_precision,*/ precision;
  precision.setZero();

  for(itctx_.Level = cfg.FirstLevel; itctx_.Level >= cfg.LastLevel; --itctx_.Level)
  {
    result.Statistics.Levels.push_back(LevelStats());
    LevelStats& level_stats = result.Statistics.Levels.back();

    mean.setZero();
    precision.setZero();

    // reset error after every pyramid level? yes because errors from different levels are not comparable
    itctx_.Iteration = 0;
    itctx_.Error = std::numeric_limits<double>::max();

    RgbdImage& cur = current.level(itctx_.Level);
    const IntrinsicMatrix& K = cur.camera().intrinsics();

    Vector8f wcur, wref;
    // i z idx idy zdx zdy
    float wcur_id = 0.5f, wref_id = 0.5f, wcur_zd = 1.0f, wref_zd = 0.0f;

    wcur <<  1.0f / 255.0f,  1.0f, wcur_id * K.fx() / 255.0f, wcur_id * K.fy() / 255.0f, wcur_zd * K.fx(), wcur_zd * K.fy(), 0.0f, 0.0f;
    wref << -1.0f / 255.0f, -1.0f, wref_id * K.fx() / 255.0f, wref_id * K.fy() / 255.0f, wref_zd * K.fx(), wref_zd * K.fy(), 0.0f, 0.0f;

//    sw_prep[itctx_.Level].start();

    PointSelection::PointIterator first_point, last_point;
    reference.select(itctx_.Level, first_point, last_point);
    cur.buildAccelerationStructure();

    level_stats.Id = itctx_.Level;
    level_stats.MaxValidPixels = reference.getMaximumNumberOfPoints(itctx_.Level);
    level_stats.ValidPixels = last_point - first_point;

//    sw_prep[itctx_.Level].stopAndPrint();

    NormalEquationsLeastSquares ls;
    Matrix6d A;
    Vector6d x, b;
    x = inc.log();

    ComputeResidualsResult compute_residuals_result;
    compute_residuals_result.first_point_error = points_error.begin();
    compute_residuals_result.first_residual = residuals.begin();
    compute_residuals_result.first_valid_flag = valid_residuals.begin();

//    sw_level[itctx_.Level].start();
    do
    {
      level_stats.Iterations.push_back(IterationStats());
      IterationStats& iteration_stats = level_stats.Iterations.back();
      iteration_stats.Id = itctx_.Iteration;

//      sw_it[itctx_.Level].start();

      double total_error = 0.0f;
//      sw_error[itctx_.Level].start();
      Eigen::Affine3f transformf;

        inc = Sophus::SE3d::exp(x);
        initial.update() = inc.inverse() * initial();
        estimate.update() = inc * estimate();

        transformf = estimate().matrix().cast<float>();

        if(debug)
        {
          dvo::core::computeResidualsAndValidFlagsSse(first_point, last_point, cur, K, transformf, wref, wcur, compute_residuals_result);
        }
        else
        {
          dvo::core::computeResidualsSse(first_point, last_point, cur, K, transformf, wref, wcur, compute_residuals_result);
        }
        size_t n = (compute_residuals_result.last_residual - compute_residuals_result.first_residual);
        iteration_stats.ValidConstraints = n;

        if(n < 6)
        {
          initial.revert();
          estimate.revert();

          level_stats.TerminationCriterion = TerminationCriteria::TooFewConstraints;

          break;
        }

        if(itctx_.IsFirstIterationOnLevel())
        {
          std::fill(weights.begin(), weights.begin() + n, 1.0f);
        }
        else
        {
          dvo::core::computeWeightsSse(compute_residuals_result.first_residual, compute_residuals_result.last_residual, weights.begin(), mean, precision);
        }

        precision = dvo::core::computeScaleSse(compute_residuals_result.first_residual, compute_residuals_result.last_residual, weights.begin(), mean).inverse();

        float ll = computeCompleteDataLogLikelihood(compute_residuals_result.first_residual, compute_residuals_result.last_residual, weights.begin(), mean, precision);

        iteration_stats.TDistributionLogLikelihood = -ll;
        iteration_stats.TDistributionMean = mean.cast<double>();
        iteration_stats.TDistributionPrecision = precision.cast<double>();
        iteration_stats.PriorLogLikelihood = cfg.Mu * initial().log().squaredNorm();

        total_error = -ll;//iteration_stats.TDistributionLogLikelihood + iteration_stats.PriorLogLikelihood;

        itctx_.LastError = itctx_.Error;
        itctx_.Error = total_error;

//      sw_error[itctx_.Level].stopAndPrint();

      // accept the last increment?
      accept = itctx_.Error < itctx_.LastError;

      if(!accept)
      {
        initial.revert();
        estimate.revert();

        level_stats.TerminationCriterion = TerminationCriteria::LogLikelihoodDecreased;

        break;
      }

      // now build equation system
//      sw_linsys[itctx_.Level].start();

      WeightVectorType::iterator w_it = weights.begin();

      Matrix2x6 J, Jw;
      Eigen::Vector2f Ji;
      Vector6 Jz;
      ls.initialize(1);
      for(PointIterator e_it = compute_residuals_result.first_point_error; e_it != compute_residuals_result.last_point_error; ++e_it, ++w_it)
      {
        computeJacobianOfProjectionAndTransformation(e_it->getPointVec4f(), Jw);
        compute3rdRowOfJacobianOfTransformation(e_it->getPointVec4f(), Jz);

        J.row(0) = e_it->getIntensityDerivativeVec2f().transpose() * Jw;
        J.row(1) = e_it->getDepthDerivativeVec2f().transpose() * Jw - Jz.transpose();

        ls.update(J, e_it->getIntensityAndDepthVec2f(), (*w_it) * precision);
      }
      ls.finish();

      A = ls.A.cast<double>() + cfg.Mu * Matrix6d::Identity();
      b = ls.b.cast<double>() + cfg.Mu * initial().log();
      x = A.ldlt().solve(b);

//      sw_linsys[itctx_.Level].stopAndPrint();

      iteration_stats.EstimateIncrement = x;
      iteration_stats.EstimateInformation = A;

      itctx_.Iteration++;
//      sw_it[itctx_.Level].stopAndPrint();
    }
    while(accept && x.lpNorm<Eigen::Infinity>() > cfg.Precision && !itctx_.IterationsExceeded());

    if(x.lpNorm<Eigen::Infinity>() <= cfg.Precision)
      level_stats.TerminationCriterion = TerminationCriteria::IncrementTooSmall;

    if(itctx_.IterationsExceeded())
      level_stats.TerminationCriterion = TerminationCriteria::IterationsExceeded;

//    sw_level[itctx_.Level].stopAndPrint();
  }

  LevelStats& last_level = result.Statistics.Levels.back();
  IterationStats& last_iteration = last_level.TerminationCriterion != TerminationCriteria::LogLikelihoodDecreased ? last_level.Iterations[last_level.Iterations.size() - 1] : last_level.Iterations[last_level.Iterations.size() - 2];

  result.Transformation = estimate().inverse().matrix();
  result.Information = last_iteration.EstimateInformation * 0.008 * 0.008;
  result.LogLikelihood = last_iteration.TDistributionLogLikelihood + last_iteration.PriorLogLikelihood;

  return success;
}

cv::Mat DenseTracker::computeIntensityErrorImage(dvo::core::RgbdImagePyramid& reference, dvo::core::RgbdImagePyramid& current, const dvo::core::AffineTransformd& transformation, size_t level)
{
  reference.compute(level + 1);
  current.compute(level + 1);
  reference_selection_.setRgbdImagePyramid(reference);
  reference_selection_.debug(true);

  std::vector<uint8_t> valid_residuals;

  if(points_error.size() < reference_selection_.getMaximumNumberOfPoints(level))
    points_error.resize(reference_selection_.getMaximumNumberOfPoints(level));
  if(residuals.size() < reference_selection_.getMaximumNumberOfPoints(level))
    residuals.resize(reference_selection_.getMaximumNumberOfPoints(level));

  valid_residuals.resize(reference_selection_.getMaximumNumberOfPoints(level));

  PointSelection::PointIterator first_point, last_point;
  reference_selection_.select(level, first_point, last_point);

  RgbdImage& cur = current.level(level);
  cur.buildAccelerationStructure();
  const IntrinsicMatrix& K = cur.camera().intrinsics();

  Vector8f wcur, wref;
  // i z idx idy zdx zdy
  float wcur_id = 0.5f, wref_id = 0.5f, wcur_zd = 1.0f, wref_zd = 0.0f;

  wcur <<  1.0f / 255.0f,  1.0f, wcur_id * K.fx() / 255.0f, wcur_id * K.fy() / 255.0f, wcur_zd * K.fx(), wcur_zd * K.fy(), 0.0f, 0.0f;
  wref << -1.0f / 255.0f, -1.0f, wref_id * K.fx() / 255.0f, wref_id * K.fy() / 255.0f, wref_zd * K.fx(), wref_zd * K.fy(), 0.0f, 0.0f;

  ComputeResidualsResult compute_residuals_result;
  compute_residuals_result.first_point_error = points_error.begin();
  compute_residuals_result.first_residual = residuals.begin();
  compute_residuals_result.first_valid_flag = valid_residuals.begin();

  dvo::core::computeResidualsAndValidFlagsSse(first_point, last_point, cur, K, transformation.cast<float>(), wref, wcur, compute_residuals_result);

  cv::Mat result = cv::Mat::zeros(reference.level(level).intensity.size(), CV_32FC1), debug_idx;

  reference_selection_.getDebugIndex(level, debug_idx);

  uint8_t *valid_pixel_it = debug_idx.ptr<uint8_t>();
  ValidFlagIterator valid_residual_it = compute_residuals_result.first_valid_flag;
  ResidualIterator residual_it = compute_residuals_result.first_residual;

  float *result_it = result.ptr<float>();
  float *result_end = result_it + result.total();

  for(; result_it != result_end; ++result_it)
  {
    if(*valid_pixel_it == 1)
    {
      if(*valid_residual_it == 1)
      {
        *result_it = std::abs(residual_it->coeff(0));

        ++residual_it;
      }
      ++valid_residual_it;
    }
    ++valid_pixel_it;
  }

  reference_selection_.debug(false);

  return result;
}

// jacobian computation
inline void DenseTracker::computeJacobianOfProjectionAndTransformation(const Vector4& p, Matrix2x6& j)
{
  NumType z = 1.0f / p(2);
  NumType z_sqr = 1.0f / (p(2) * p(2));

  j(0, 0) =  z;
  j(0, 1) =  0.0f;
  j(0, 2) = -p(0) * z_sqr;
  j(0, 3) = j(0, 2) * p(1);//j(0, 3) = -p(0) * p(1) * z_sqr;
  j(0, 4) = 1.0f - j(0, 2) * p(0);//j(0, 4) =  (1.0 + p(0) * p(0) * z_sqr);
  j(0, 5) = -p(1) * z;

  j(1, 0) =  0.0f;
  j(1, 1) =  z;
  j(1, 2) = -p(1) * z_sqr;
  j(1, 3) = -1.0f + j(1, 2) * p(1); //j(1, 3) = -(1.0 + p(1) * p(1) * z_sqr);
  j(1, 4) = -j(0, 3); //j(1, 4) =  p(0) * p(1) * z_sqr;
  j(1, 5) =  p(0) * z;
}

inline void DenseTracker::compute3rdRowOfJacobianOfTransformation(const Vector4& p, Vector6& j)
{
  j(0) = 0.0;
  j(1) = 0.0;
  j(2) = 1.0;
  j(3) = p(1);
  j(4) = -p(0);
  j(5) = 0.0;
}

} /* namespace dvo */


// ====================================================================
// This file is part of FlexibleSUSY.
//
// FlexibleSUSY 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.
//
// FlexibleSUSY 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 FlexibleSUSY.  If not, see
// <http://www.gnu.org/licenses/>.
// ====================================================================

// File generated at Sat 27 Aug 2016 12:47:55

#ifndef NUHMSSM_SLHA_IO_H
#define NUHMSSM_SLHA_IO_H

#include "NUHMSSM_two_scale_model_slha.hpp"
#include "NUHMSSM_info.hpp"
#include "NUHMSSM_observables.hpp"
#include "NUHMSSM_physical.hpp"
#include "slha_io.hpp"
#include "ckm.hpp"
#include "ew_input.hpp"
#include "lowe.h"

#include <Eigen/Core>
#include <string>
#include <utility>

#define Pole(p) physical.p
#define PHYSICAL(p) model.get_physical().p
#define PHYSICAL_SLHA(p) model.get_physical_slha().p
#define LOCALPHYSICAL(p) physical.p
#define MODEL model
#define MODELPARAMETER(p) model.get_##p()
#define OBSERVABLES observables
#define LowEnergyConstant(p) Electroweak_constants::p
#define SCALES(p) scales.p

namespace flexiblesusy {

struct NUHMSSM_input_parameters;
class Spectrum_generator_settings;

struct NUHMSSM_scales {
   NUHMSSM_scales() : HighScale(0.), SUSYScale(0.), LowScale(0.) {}
   double HighScale, SUSYScale, LowScale;
};

class NUHMSSM_slha_io {
public:
   NUHMSSM_slha_io();
   ~NUHMSSM_slha_io() {}

   void clear();

   void fill(softsusy::QedQcd& qedqcd) const { slha_io.fill(qedqcd); }
   void fill(NUHMSSM_input_parameters&) const;
   void fill(NUHMSSM_mass_eigenstates&) const;
   template <class T> void fill(NUHMSSM_slha<T>&) const;
   void fill(Physical_input&) const;
   void fill(Spectrum_generator_settings&) const;
   double get_parameter_output_scale() const;
   const SLHA_io& get_slha_io() const { return slha_io; }
   void read_from_file(const std::string&);
   void read_from_source(const std::string&);
   void read_from_stream(std::istream&);
   void set_block(const std::string& str, SLHA_io::Position position = SLHA_io::back) { slha_io.set_block(str, position); }
   void set_blocks(const std::vector<std::string>& vec, SLHA_io::Position position = SLHA_io::back) { slha_io.set_blocks(vec, position); }
   void set_extpar(const NUHMSSM_input_parameters&);
   template <class T> void set_extra(const NUHMSSM_slha<T>&, const NUHMSSM_scales&, const NUHMSSM_observables&);
   void set_minpar(const NUHMSSM_input_parameters&);
   void set_sminputs(const softsusy::QedQcd&);
   template <class T> void set_spectrum(const NUHMSSM_slha<T>&);
   template <class T> void set_spectrum(const NUHMSSM<T>&);
   void set_spinfo(const Problems<NUHMSSM_info::NUMBER_OF_PARTICLES>&);
   void set_print_imaginary_parts_of_majorana_mixings(bool);
   void write_to_file(const std::string&);
   void write_to_stream(std::ostream& ostr = std::cout) { slha_io.write_to_stream(ostr); }

   static void fill_minpar_tuple(NUHMSSM_input_parameters&, int, double);
   static void fill_extpar_tuple(NUHMSSM_input_parameters&, int, double);

   template <class T>
   static void fill_slhaea(SLHAea::Coll&, const NUHMSSM_slha<T>&, const softsusy::QedQcd&, const NUHMSSM_scales&, const NUHMSSM_observables&);

   template <class T>
   static SLHAea::Coll fill_slhaea(const NUHMSSM_slha<T>&, const softsusy::QedQcd&, const NUHMSSM_scales&, const NUHMSSM_observables&);

private:
   SLHA_io slha_io; ///< SLHA io class
   bool print_imaginary_parts_of_majorana_mixings;
   static unsigned const NUMBER_OF_DRBAR_BLOCKS = 14;
   static char const * const drbar_blocks[NUMBER_OF_DRBAR_BLOCKS];

   void set_mass(const NUHMSSM_physical&, bool);
   void set_mixing_matrices(const NUHMSSM_physical&, bool);
   template <class T> void set_model_parameters(const NUHMSSM_slha<T>&);
   void set_ckm(const Eigen::Matrix<std::complex<double>,3,3>&, double);
   void set_pmns(const Eigen::Matrix<std::complex<double>,3,3>&, double);
   double read_scale() const;
   void fill_drbar_parameters(NUHMSSM_mass_eigenstates&) const;
   void fill_physical(NUHMSSM_physical&) const;
};

/**
 * Reads DR-bar parameters, pole masses and mixing matrices from a
 * SLHA output file.
 */
template <class T>
void NUHMSSM_slha_io::fill(NUHMSSM_slha<T>& model) const
{
   fill(static_cast<NUHMSSM_mass_eigenstates&>(model));
   fill_physical(model.get_physical_slha());
}

template <class T>
void NUHMSSM_slha_io::fill_slhaea(
   SLHAea::Coll& slhaea, const NUHMSSM_slha<T>& model,
   const softsusy::QedQcd& qedqcd, const NUHMSSM_scales& scales,
   const NUHMSSM_observables& observables)
{
   NUHMSSM_slha_io slha_io;
   const NUHMSSM_input_parameters& input = model.get_input();
   const Problems<NUHMSSM_info::NUMBER_OF_PARTICLES>& problems
      = model.get_problems();
   const bool error = problems.have_problem();

   slha_io.set_spinfo(problems);
   slha_io.set_sminputs(qedqcd);
   slha_io.set_minpar(input);
   slha_io.set_extpar(input);
   if (!error) {
      slha_io.set_spectrum(model);
      slha_io.set_extra(model, scales, observables);
   }

   slhaea = slha_io.get_slha_io().get_data();
}

template <class T>
SLHAea::Coll NUHMSSM_slha_io::fill_slhaea(
   const NUHMSSM_slha<T>& model, const softsusy::QedQcd& qedqcd,
   const NUHMSSM_scales& scales, const NUHMSSM_observables& observables)
{
   SLHAea::Coll slhaea;
   NUHMSSM_slha_io::fill_slhaea(slhaea, model, qedqcd, scales, observables);

   return slhaea;
}

/**
 * Stores the model (DR-bar) parameters in the SLHA object.
 *
 * @param model model class
 */
template <class T>
void NUHMSSM_slha_io::set_model_parameters(const NUHMSSM_slha<T>& model)
{
   {
      std::ostringstream block;
      block << "Block gauge Q= " << FORMAT_SCALE(model.get_scale()) << '\n'
            << FORMAT_ELEMENT(1, (MODELPARAMETER(g1) * 0.7745966692414834), "gY")
            << FORMAT_ELEMENT(2, (MODELPARAMETER(g2)), "g2")
            << FORMAT_ELEMENT(3, (MODELPARAMETER(g3)), "g3")
      ;
      slha_io.set_block(block);
   }
   slha_io.set_block("Yu", ToMatrix(MODELPARAMETER(Yu_slha)), "Yu", model.get_scale());
   slha_io.set_block("Yd", ToMatrix(MODELPARAMETER(Yd_slha)), "Yd", model.get_scale());
   slha_io.set_block("Ye", ToMatrix(MODELPARAMETER(Ye_slha)), "Ye", model.get_scale());
   slha_io.set_block("Te", MODELPARAMETER(TYe_slha), "TYe", model.get_scale());
   slha_io.set_block("Td", MODELPARAMETER(TYd_slha), "TYd", model.get_scale());
   slha_io.set_block("Tu", MODELPARAMETER(TYu_slha), "TYu", model.get_scale());
   {
      std::ostringstream block;
      block << "Block HMIX Q= " << FORMAT_SCALE(model.get_scale()) << '\n'
            << FORMAT_ELEMENT(1, (MODELPARAMETER(Mu)), "Mu")
            << FORMAT_ELEMENT(101, (MODELPARAMETER(BMu)), "BMu")
            << FORMAT_ELEMENT(102, (MODELPARAMETER(vd)), "vd")
            << FORMAT_ELEMENT(103, (MODELPARAMETER(vu)), "vu")
      ;
      slha_io.set_block(block);
   }
   slha_io.set_block("MSQ2", MODELPARAMETER(mq2_slha), "mq2", model.get_scale());
   slha_io.set_block("MSE2", MODELPARAMETER(me2_slha), "me2", model.get_scale());
   slha_io.set_block("MSL2", MODELPARAMETER(ml2_slha), "ml2", model.get_scale());
   slha_io.set_block("MSU2", MODELPARAMETER(mu2_slha), "mu2", model.get_scale());
   slha_io.set_block("MSD2", MODELPARAMETER(md2_slha), "md2", model.get_scale());
   {
      std::ostringstream block;
      block << "Block MSOFT Q= " << FORMAT_SCALE(model.get_scale()) << '\n'
            << FORMAT_ELEMENT(21, (MODELPARAMETER(mHd2)), "mHd2")
            << FORMAT_ELEMENT(22, (MODELPARAMETER(mHu2)), "mHu2")
            << FORMAT_ELEMENT(1, (MODELPARAMETER(MassB)), "MassB")
            << FORMAT_ELEMENT(2, (MODELPARAMETER(MassWB)), "MassWB")
            << FORMAT_ELEMENT(3, (MODELPARAMETER(MassG)), "MassG")
      ;
      slha_io.set_block(block);
   }

   {
      std::ostringstream block;
      block << "Block Phases Q= " << FORMAT_SCALE(model.get_scale()) << '\n'
            << FORMAT_ELEMENT(1, (Re(MODELPARAMETER(PhaseGlu))), "Re(PhaseGlu)")
      ;
      slha_io.set_block(block);
   }
   {
      std::ostringstream block;
      block << "Block IMPhases Q= " << FORMAT_SCALE(model.get_scale()) << '\n'
            << FORMAT_ELEMENT(1, (Im(MODELPARAMETER(PhaseGlu))), "Im(PhaseGlu)")
      ;
      slha_io.set_block(block);
   }

}

/**
 * Writes extra SLHA blocks
 *
 * @param model model class
 */
template <class T>
void NUHMSSM_slha_io::set_extra(
   const NUHMSSM_slha<T>& model, const NUHMSSM_scales& scales,
   const NUHMSSM_observables& observables)
{
   const NUHMSSM_physical physical(model.get_physical_slha());

}

/**
 * Stores the model (DR-bar) parameters, masses and mixing matrices in
 * the SLHA object.
 *
 * @param model model class in BPMZ convention
 */
template <class T>
void NUHMSSM_slha_io::set_spectrum(const NUHMSSM<T>& model)
{
   const NUHMSSM_slha<T> model_slha(model);
   set_spectrum(model_slha);
}

/**
 * Stores the model (DR-bar) parameters, masses and mixing matrices in
 * the SLHA object.
 *
 * @param model model class in SLHA convention
 */
template <class T>
void NUHMSSM_slha_io::set_spectrum(const NUHMSSM_slha<T>& model)
{
   const NUHMSSM_physical physical(model.get_physical_slha());
   const bool write_sm_masses = model.do_calculate_sm_pole_masses();

   set_model_parameters(model);
   set_mass(physical, write_sm_masses);
   set_mixing_matrices(physical, write_sm_masses);

   if (slha_io.get_modsel().quark_flavour_violated)
      set_ckm(model.get_ckm_matrix(), model.get_scale());

   if (slha_io.get_modsel().lepton_flavour_violated)
      set_pmns(model.get_pmns_matrix(), model.get_scale());
}

} // namespace flexiblesusy

#undef Pole
#undef PHYSICAL
#undef PHYSICAL_SLHA
#undef LOCALPHYSICAL
#undef MODEL
#undef MODELPARAMETER
#undef OBSERVABLES
#undef LowEnergyConstant
#undef SCALES

#endif


/* Copyright (C) 2012-2019 IBM Corp.
 * This program is Licensed under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance
 * with the License. You may obtain a copy of the License at
 *   http://www.apache.org/licenses/LICENSE-2.0
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License. See accompanying LICENSE file.
 */
#include <NTL/ZZ.h>
#include "FHEContext.h"
#include "Ctxt.h"
#include "permutations.h"
#include "EncryptedArray.h"

namespace helib {

std::ostream& operator<< (std::ostream &s, const PermNetwork &net)
{
  s << "[";
  for (long i=0; i< net.layers.length(); i++) {
    const PermNetLayer& lyr = net.layers[i];
    s << "[" << lyr.genIdx << " " << lyr.e << " " << lyr.isID << " "
      << lyr.shifts << "]\n";
  }
  return s << "]";
}

// Compute one or more layers corresponding to one network of one leaf
void PermNetwork::setLayers4Leaf(long lyrIdx, const ColPerm& p,
				 const NTL::Vec<long>& benesLvls, long gIdx,
				 const SubDimension& leafData,
				 const Permut& map2cube)
{
#ifdef DEBUG_PRINTOUT
  std::cerr << "Layer "<<lyrIdx<<", column-permutation="<< p << std::endl;
#endif
  // Compute the shift amounts for all the layers in this network
  NTL::Vec<bool> isID;
  NTL::Vec<Permut> shifts;
  if (benesLvls.length()==1) {// Special case for a "trivial" 1-layer network
    shifts.SetLength(1);
    isID.SetLength(1);
    isID[0] = !p.getShiftAmounts(shifts[0]);
  }
  else  // The general case of a multi-layer Benes network
    p.getBenesShiftAmounts(shifts,isID,benesLvls);

  // Copy the shift amounts to the right place in the bigger network,
  // renaming the slots from a linear array to the hyper cube
  for (long i=0; i<benesLvls.length(); i++) {
    PermNetLayer& lyr = layers[lyrIdx+i];
    lyr.genIdx = gIdx;
    lyr.isID = isID[i];
    lyr.e = leafData.e;
    if (!lyr.isID) {
#ifdef DEBUG_PRINTOUT
      std::cerr << "layer "<<lyrIdx+i<<": "<<shifts[i]<<std::endl;
#endif
      if (leafData.good) // For good leaves, shift by -x is the same as size-x
	for (long k=0; k<shifts[i].length(); k++)
	  if (shifts[i][k]<0) shifts[i][k] += leafData.size;
      applyPermToVec(lyr.shifts, shifts[i], map2cube); // do the renaming
#ifdef DEBUG_PRINTOUT
      std::cerr << "       : "<<lyr.shifts<<std::endl;
#endif
    }
    //    else std::cerr << "layer "<<lyrIdx+i<<"= identity\n";
  }
}

// Build a full permutation network
void PermNetwork::buildNetwork(const Permut& pi, const GeneratorTrees& trees)
{
  if (trees.numTrees()==0) { // the identity permutation, nothing to do
    layers.SetLength(0);
    return;
  }

  NTL::Vec<long> dims;
  trees.getCubeSubDims(dims);

  //  std::cerr << "pi =      "<<pi<<std::endl;
  //  std::cerr << "map2cube ="<<trees.mapToCube()<<std::endl;
  //  std::cerr << "map2array="<<trees.mapToArray()<<std::endl;

  // Compute the permutation on the cube, rho = map2cube o pi o map2array

  Permut rho;
  applyPermsToVec(rho, trees.mapToCube(), pi, trees.mapToArray());
  //  std::cerr << "rho =     "<<rho<<std::endl;

  // Break rho along the different dimensions
  CubeSignature sig(dims); // make a cube-signature object
  std::vector<ColPerm> perms;
  breakPermByDim(perms, rho, sig);

  //  for (long i=0; i<(long)perms.size(); i++) { // debugging printouts
  //    Permut tmp;
  //    perms[i].makeExplicit(tmp);
  //    std::cerr << " prems["<<i<<"]="<<tmp<<std::endl;
  //  }

  layers.SetLength(trees.numLayers()); // allocate space

  // Go over the different permutations and build the corresponding layers
  long dimIdx =0;
  long frntLyr=0, backLyr=layers.length();
  for (long g=0; g<trees.numTrees(); g++) { // go over all the generators/trees
    const OneGeneratorTree &T = trees[g];

    // In each tree, go over all the leaves
    for (long leaf=T.firstLeaf(); leaf>=0; leaf=T.nextLeaf(leaf)) {
      const SubDimension& leafData = T[leaf].getData();

      // This leaf determines layers frntLyer...frntLey+frst.length()-1, and
      // if it isn't the middle then also backLyr-scnd.length()...backLyr-1

      // handle the first Benes network
      setLayers4Leaf(/*1st-layer-index=*/frntLyr,
		     /*permutation    =*/perms[dimIdx],
		     /*Benes levels   =*/leafData.frstBenes,
		     /*generator index=*/T.getAuxKey(),
		     /*(size,good,e)  =*/leafData,
		     /*hypercube renaming permutation=*/trees.mapToCube());
      frntLyr += leafData.frstBenes.length(); // how many layers were used
      dimIdx++;

      if (leafData.scndBenes.length()>0) { // Also a second Benes network
	long dimIdx2 = perms.size() -dimIdx; // dimIdx was incremented above
	backLyr -= leafData.scndBenes.length();
	setLayers4Leaf(/*1st-layer-index=*/backLyr,
		       /*permutation    =*/perms[dimIdx2],
		       /*Benes levels   =*/leafData.scndBenes,
		       /*generator index=*/T.getAuxKey(),
		       /*(size,good,e)  =*/leafData,
		       /*hypercube renaming permutation=*/trees.mapToCube());
      }
    }
  }
}

// Apply a permutation network to a hypercube, used mostly for debugging
void PermNetwork::applyToCube(HyperCube<long>& cube) const
{
  if (layers.length()==0) return;
  long n = cube.getSize();
  NTL::Vec<long> tmp(NTL::INIT_SIZE, n); // temporary vector

  // Apply the layers, one at a time
  for (long i=0; i<layers.length(); i++) {
    const PermNetLayer& lyr = layers[i];
    if (lyr.isID) continue; // this layer is the identity permutation

    //OLD: assert(lyr.shifts.length()==n);
    helib::assertEq(lyr.shifts.length(), n, "layer has incorrect size");

    // This layer shift elements along the dimension lyr.genIdx
    long dim = lyr.genIdx;

    // Move elements as dictated by this layer
    for (long j=0; j<n; j++) {
      long shamt = lyr.e * lyr.shifts[j]; // how much to shift this slot
      if (shamt<0) shamt += cube.getDim(dim); // addCoord expects shamt>=0
      long j2 = cube.addCoord(j, dim, shamt); // new index for this slot
      tmp[j2] = cube[j];
    }
    // Copy back to cube
    for (long j=0; j<n; j++)
      cube[j] = tmp[j];
#ifdef DEBUG_PRINTOUT
    std::cerr << " after layer "<< i << ", cube=" << cube.getData()<<std::endl;
#endif
  }
}

void PermNetwork::applyToPtxt(NTL::ZZX& p, const EncryptedArray& ea) const
{
  throw helib::LogicError("PermNetwork::applyToPtxt is not implemented");
}

// Upon return, mask[i]=1 if haystack[i]=needle, 0 otherwise.
// Also set to 0 all the entries in haystack where mask[i]=1.
// Return the index of the first nonzero entry in haystack at the end
// of the pass (-1 if they are all zero). Also return a flag saying if
// any entries of the mask are nonzero.
static std::pair<long,bool>
makeMask(std::vector<long>& mask, NTL::Vec<long>& haystack, long needle)
{
  long found = false;
  long fstNonZeroIdx = -1;
  for (long i=0; i<(long)mask.size(); i++) {
    if (haystack[i] == needle) { // found a needle
      found = true;
      mask[i]=1;
      haystack[i]=0; // remove this needle from haystack
    } else {         // no needle here
      mask[i]=0;
      if (haystack[i]!=0 && fstNonZeroIdx<0)
	fstNonZeroIdx = i;       // first nonzero entry in haystack
    }
  }
  return std::make_pair(fstNonZeroIdx,found);
}

// Apply a permutation network to a ciphertext
void PermNetwork::applyToCtxt(Ctxt& c, const EncryptedArray& ea) const
{
  const PAlgebra& al = ea.getPAlgebra();

  // Apply the layers, one at a time
  for (long i=0; i<layers.length(); i++) {
    const PermNetLayer& lyr = layers[i];
    if (lyr.isID) continue; // this layer is the identity permutation

    // This layer is shifted via powers of g^e mod m
    long g2e = NTL::PowerMod(al.ZmStarGen(lyr.genIdx), lyr.e, al.getM());

    NTL::Vec<long> unused = lyr.shifts; // copy to a new vector
    std::vector<long> mask(lyr.shifts.length());  // buffer to hold masks
    Ctxt sum(c.getPubKey(), c.getPtxtSpace()); // an empty ciphertext

    long shamt = 0;
    bool frst = true;
    while (true) {
      std::pair<long,bool> ret=makeMask(mask, unused, shamt); // compute mask
      if (ret.second) { // non-empty mask
	Ctxt tmp = c;
	NTL::ZZX maskPoly;
	ea.encode(maskPoly, mask);    // encode mask as polynomial
	tmp.multByConstant(maskPoly); // multiply by mask
	if (shamt!=0) // rotate if the shift amount is nonzero
	  tmp.smartAutomorph(NTL::PowerMod(g2e, shamt, al.getM()));
	if (frst) {
	  sum = tmp;
	  frst = false;
	}
	else
	  sum += tmp;
      }
      if (ret.first >= 0)
	shamt = unused[ret.first]; // next shift amount to use

      else break; // unused is all-zero, done with this layer
    }
    c = sum; // update the cipehrtext c before the next layer
  }
}

}