/*
 * Copyright (c) 2009 Carnegie Mellon University.
 *     All rights reserved.
 *
 *  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.
 *
 *
 */

/**
 *
 * \brief This application performs MAP inference on Markov Nets
 * provided in standard UAI file format via Dual-Decomposition.
 *
 *
 *  \authors Dhruv Batra, André Martins, Aroma Mahendru
 */

#ifndef __DD_GRLAB_HPP__
#define __DD_GRLAB_HPP__

#include <boost/config/warning_disable.hpp>
#include <boost/spirit/include/qi.hpp>
#include <boost/spirit/include/phoenix_core.hpp>
#include <boost/spirit/include/phoenix_operator.hpp>
#include <boost/spirit/include/phoenix_stl.hpp>

#include <Eigen/Dense>
#include "eigen_serialization.hpp"

#include "dd_opts.hpp"
#include "utils.hpp"

#include <graphlab.hpp>
#include <graphlab/macros_def.hpp>

using namespace std;

/**
 * \brief Eigen library vectors are used to store the potentials (log-space)
 */
typedef Eigen::VectorXd factor_type;
typedef Eigen::VectorXd vec;
typedef Eigen::MatrixXd mat;

struct dd_global_vars {

double old_dual ;        // stores the value of dual objective for the previous iteration
double primal_best ;     //  stores the value of best primal objective found so far.
bool converged ;         // true if dual objective value has converged to required tolerance level, otherwise false
int dual_inc_count ;     // keeps track of the number of times the value of dual objective increased
vector < vector<double> > history ; // stores dual and primal objective values
int sq_norm_g ;   //  stores the value of the square of the norm of the subgradient vector
int iter_at_aggregate ;  //  iteration number at the time of aggregate
graphlab::timer timer ; //  time object. Helps in finding the time elapsed.
dd_global_vars(): old_dual(200), primal_best(-1e10),
                  converged(false), dual_inc_count(1),
                  history(4,vector<double>()),
                  sq_norm_g(100),
                  iter_at_aggregate(0){};
} global_vars;

/* end of struct dd_global_vars */

/////////////////////////////////////////////////////////////////////////
// Edge and Vertex data and Graph Type
/**
 * \brief There is a vertex for each factor in the graph AND each singleton
 */
struct vertex_data
{
    int nvars;                 // Number of variables in this factor.
    int factor_type;           //type of  factor : dense(0), budget(1)
    int degree;                // Degree of this factor (same as nvars for higher-order factors).

    vector<int> cards;         // Cardinality of each variable.
    vector<int> neighbors;     // Vertex ids of the neighbors.
    vec potentials;            // Potentials for each configuration of the factor.
    int budget;                // Only for Budget factors
    vector <int> bound_states;  // Only for Budget Factors
    vector<int> unary_degree;   // Only for unary vertices

    int best_configuration;    // Index of the best configuration at a subgradient step.
                               // TODO: Maybe replace best_configuration by beliefs for the high order variables?
                               // In which case, beliefs would be vector<vec> beliefs.

    double dual_contrib;       // Contribution of this factor to the dual. We can compute this during the applies.
    double primal_contrib;     // Contribution of this factor to the primal. We can compute this during the applies.
                               // NOTENOTE: ONLY true for sync engine. For async, we need to write an aggregator function.
    double dual_res_contrib;   // Contribution of this factor to the dual residual. (Used only for ADMM)
    double primal_res_contrib; // Contribution of this factor to the primal residual. (Used only for ADMM)
    double primal_rel_contrib; // Contribution of this factor to the relaxed primal. (Used only for ADMM)

    vec beliefs;               // Posterior values for the configurations after averaging (projected DD, unary variables only).
    vec factor_beliefs;        // Posterior value for factor variables
    int apply_count;           // No. of times apply has been called on this vertex
    int sum_sq_norm_g;         // Sum of square of norm of subgradient for each vertex (used only for factor vertices)

    bool schedule_vertex;      // Decides if vertex is to be scheduled for further iterations or not

    vector<double> distribution_;
    vector<int> active_set_;
    vector<double> inverse_A_;
    vector<pair<double, int> > last_sort_;

    vertex_data():
    nvars(0), factor_type(0), degree(0),
    budget(0),
    best_configuration(0),
    dual_contrib(0), primal_contrib(0),
    dual_res_contrib(0), primal_res_contrib(0), primal_rel_contrib(0),
    sum_sq_norm_g(0),
    apply_count(0),
    schedule_vertex(true)
    {}

    void load(graphlab::iarchive& arc)
    {
        arc >> nvars >> degree
            >> cards >> neighbors >> potentials
            >> dual_contrib >> primal_contrib
            >> best_configuration >> beliefs
            >> apply_count>>factor_beliefs
            >>sum_sq_norm_g>>primal_res_contrib
            >>primal_rel_contrib>>dual_res_contrib
            >>schedule_vertex>>factor_type
            >>budget>>bound_states
            >>unary_degree>>active_set_
            >>inverse_A_>>last_sort_ ;
    }
    void save(graphlab::oarchive& arc) const
    {
        arc << nvars << degree
            << cards << neighbors << potentials
            << dual_contrib << primal_contrib
            << best_configuration << beliefs
            << apply_count <<factor_beliefs
            << sum_sq_norm_g<<primal_res_contrib
            <<primal_rel_contrib<<dual_res_contrib
            <<schedule_vertex<<factor_type
            <<budget<<bound_states
            <<unary_degree<<active_set_
            <<inverse_A_<<last_sort_;
    }
}; // end of vertex_data

/**
 * \brief There is an edge connecting each factor to each singleton
 * in its scope.
 */
struct edge_data
{
    vec potentials;

    vec multiplier_messages; // Dual variables, i.e. Lagrangian multipliers.
    vec local_messages;      // Local MAP variables (for projected DD).

    void load(graphlab::iarchive& arc) {
        arc >> potentials >> multiplier_messages >> local_messages;
    }
    void save(graphlab::oarchive& arc) const {
        arc << potentials << multiplier_messages << local_messages;
    }
};  //end of edge_data

/**
 * \brief gather_type is a structure that will be used as the return type of gather function. It includes
 *  messages (used both for unary and factor vertices), neighbor_best_conf, neighbor_distribution (used only for
 *  factor vertices) and sq_norm_g (for storing square of norm of subgradient) for each edge.
 */

struct gather_type
{ factor_type messages;
  factor_type multipliers;
  vector <int> neighbor_conf;
  vec neighbor_distribution;
  int sq_norm_g;

    gather_type():sq_norm_g(0){};

    gather_type(factor_type f, vector <int> nc = vector <int> (), int sg = 0,
                             vec nd = vec() ): messages(f), neighbor_conf(nc),
                                   neighbor_distribution(nd), sq_norm_g(sg){};
    void load(graphlab::iarchive& arc) {
        arc >>messages>>neighbor_conf
            >>sq_norm_g>>neighbor_distribution
            >>multipliers;
    }
    void save(graphlab::oarchive& arc) const {
        arc <<messages<<neighbor_conf
            <<sq_norm_g<<neighbor_distribution
            <<multipliers;
    }

  gather_type& operator+=(const gather_type& other)
 { messages += other.messages;
   neighbor_conf += other.neighbor_conf;
   neighbor_distribution += other.neighbor_distribution;
   sq_norm_g += other.sq_norm_g;
   multipliers += other.multipliers;

   return *this;
 }

}; // end of gather_type struct

/**
 * \brief objective is a structure that is used as the summable data type for computing
 * dual, primal objectives, residuals etc with aggregator map and reduce functions.
 */

struct objective
{ double primal, dual, primal_rel, sum_sq_norm_g, total_confs, dual_res, primal_res;

objective(): primal(0), dual(0), primal_rel(0), sum_sq_norm_g(0), total_confs(1),dual_res(1), primal_res(1){};

void load(graphlab::iarchive& arc) {
        arc >>dual>>primal>>sum_sq_norm_g
            >>primal_rel>>total_confs
            >>dual_res>>primal_res;
    }
    void save(graphlab::oarchive& arc) const {
        arc <<dual<<primal<<sum_sq_norm_g
            <<primal_rel<<total_confs
            <<dual_res<<primal_res;
    }

objective& operator+=(const objective& other)
{ primal += other.primal;
   dual += other.dual;
   sum_sq_norm_g += other.sum_sq_norm_g;
   primal_rel += other.primal_rel;
   total_confs += other.total_confs;
   dual_res += other.dual_res;
   primal_res += other.primal_res;
   return *this;
 }
}; // end of structure objective

/**
 * The graph type
 */
typedef graphlab::distributed_graph<vertex_data, edge_data> graph_type;

/**
 * \brief The Dual Decomposition Vertex Program.
 */
struct dd_vertex_program :
public graphlab::ivertex_program< graph_type, gather_type,
graphlab::messages::sum_priority >,
public graphlab::IS_POD_TYPE
{

    /////////////////////////////////////////////////////////////////////////
    // Find the configuration index of a factor given the array of states.
    /////////////////////////////////////////////////////////////////////////
    int get_configuration_index(const graph_type::vertex_type& vertex,
                                const std::vector<int>& states) const
    {
        const vertex_data& vdata = vertex.data();
        int index = states[0];
        for (size_t i = 1; i < states.size(); ++i)
        {
            index *= vdata.cards[i];
            index += states[i];
        }
        return index;
    }

    /////////////////////////////////////////////////////////////////////////
    // Find the array of states corresponding to a factor configuration index.
    /////////////////////////////////////////////////////////////////////////
    void get_configuration_states(const graph_type::vertex_type& vertex,
                                  int index, std::vector<int>* states) const
    {
        const vertex_data& vdata = vertex.data();
        int tmp = 1;
        for (size_t i = 1; i < states->size(); ++i)
            tmp *= vdata.cards[i];

        (*states)[0] = index / tmp;
        for (size_t i = 1; i < states->size(); ++i)
        {
            index = index % tmp;
            tmp /= vdata.cards[i];
            (*states)[i] = index / tmp;
        }

    }

    ///////////////////////////////////////////////////////////
    // Updates stepsize according to different stepsize rules.
    ///////////////////////////////////////////////////////////

    double update_stepsize(icontext_type& context, int type, double old_dual, double primal_best,int norm_g_sq,
                                                                             int dual_inc_count,int iter_since_aggregate) const
   {  switch (type) {
      case 0: return opts.step_size;
              break;
      case 1: return(opts.step_size/(context.iteration()+2));
              break;
      case 2: return(2* opts.step_size *(old_dual-primal_best)/((norm_g_sq+1) * (iter_since_aggregate + dual_inc_count + 1)));
              break;
      case 3: return(opts.step_size/dual_inc_count + 1);
              break;
      case 4: if(context.iteration()+2 < 300)
              return( opts.step_size/(context.iteration()+2));
              else return(opts.step_size/(300));
                     }
   }

    /**
     * \brief Given an edge and a vertex return the other vertex along
     * that edge.
     */
    inline vertex_type get_other_vertex(edge_type& edge,
                                        const vertex_type& vertex) const
    {
        return vertex.id() == edge.source().id()? edge.target() : edge.source();
    }; // end of other_vertex

    virtual edge_dir_type gather_edges(icontext_type& context,
                                       const vertex_type& vertex) const = 0;
    virtual gather_type gather(icontext_type& context, const vertex_type& vertex,
                               edge_type& edge) const = 0;
    virtual void apply(icontext_type& context, vertex_type& vertex,
                       const gather_type& total) = 0;
    virtual edge_dir_type scatter_edges(icontext_type& context,
                                        const vertex_type& vertex) const = 0;
    virtual void scatter(icontext_type& context, const vertex_type& vertex,
                         edge_type& edge) const = 0;
}; // end of class bp_vertex_program

////////////////////////////////////////////////////////////////////////////////
// This class implements the "symmetric" version of dual decomposition described
// in:
// D. Sontag, A. Globerson, T. Jaakkola.
// Introduction to Dual Decomposition for Inference.
// Optimization for Machine Learning, editors S. Sra, S. Nowozin, and S. J.
// Wright: MIT Press, 2011
////////////////////////////////////////////////////////////////////////////////

struct dd_vertex_program_symmetric : public dd_vertex_program {
    /**
     * \brief Since the MRF is undirected we will use all edges for gather and
     * scatter
     */
    edge_dir_type gather_edges(icontext_type& context,
                               const vertex_type& vertex) const
    {   if(!opts.debug){
        return graphlab::ALL_EDGES; }
        else
        return graphlab::NO_EDGES;
    }; // end of gather_edges

    /**
     * \brief The gather function takes a vertex and an edge as inputs and outputs
     a vector of numeric values. Vectors of numeric values will later be summed
     over all edges incident in this vertex. So, if the vertex is a unary factor,
     we can just return the vector of Lagrange multipliers stored in "edge.messages".
     Otherwise (if vertex is a general factor), things are a little more tricky.
     Suppose the factor is linked to K variables, with cardinalities C_1, ..., C_K.
     Suppose this edge is with respect to the k-th variable. Then, we return a
     vector of size C_1 + ... + C_K which is zero everywhere except in the
     k-th slot, where the Lagrange multipliers in "edge.messages" will be copied
     to. This way, when the "gather sum" takes place, and since all these slots
     are disjoint, we will just get the Lagrange multipliers of all the variables.
     It also gathers the best_configuration of neighbors for the factors and norm
     of subgradient value for each edge.
     */
     gather_type gather(icontext_type& context, const vertex_type& vertex, edge_type& edge) const
     {
        if (opts.verbose > 1)
            cout << "gather begin" << endl;

        const vertex_type other_vertex = get_other_vertex(edge, vertex);
        const vertex_data& vdata = vertex.data();
        edge_data& edata = edge.data();

        if (vdata.nvars == 1)
        {
            // Unary factor.
            if (opts.verbose > 1)
                cout << "This unary factor has " << vertex.num_in_edges() <<
                " in edges and " << vertex.num_out_edges() << " out edges" << endl;

            if (opts.verbose>1)
            {
                cout << "Gather on (" << vertex.id() << "," << other_vertex.id()
                                     << ") called from " << vertex.id() << "\n";
                cout << "vdata.neighbours = " << vdata.neighbors << "\n";
                cout << "Message: " << edata.multiplier_messages << "\n---\n";
            }

            gather_type gather_data(edata.multiplier_messages);
            return gather_data;
        }
        else
        {
            // General factor.
            factor_type messages;
            messages.setZero(vdata.potentials.size());
            int offset = 0;
            int index_neighbor = -1;
            for (int k = 0; k < vdata.nvars; ++k)
            {
                int vertex_id = vdata.neighbors[k];
                if (vertex_id == other_vertex.id())
                {
                    index_neighbor = k;
                    break;
                }
                offset += vdata.cards[k];
            }
            CHECK_GE(index_neighbor, 0);
            vector <int> neighbor_conf(vdata.nvars, 0);
            neighbor_conf[index_neighbor] = other_vertex.data().best_configuration;

            for (int state = 0; state < vdata.cards[index_neighbor]; ++state)
            {
                messages[offset + state] = -edata.multiplier_messages[state];
            }

            if (opts.verbose>1)
            {
                cout << "Gather on (" << vertex.id() << "," << other_vertex.id() << ") called from " << vertex.id() << "\n";
                cout << "vdata.neighbours = " << vdata.neighbors << "\n";
                cout << "estimated offset = " << offset << "\n";
                cout << "Message: " << messages << "\n---\n";
            }
            vector<int> states(vdata.nvars, -1);
            get_configuration_states(vertex, vdata.best_configuration, &states);
            int sq_norm_g = (states[index_neighbor] == other_vertex.data().best_configuration)?0:2;
            gather_type gather_data(messages, neighbor_conf);
            gather_data.sq_norm_g = sq_norm_g;
            return gather_data;
        }
        if (opts.verbose > 2)
            cout << "gather end" << endl;
    }; // end of gather function

    /**
     * \brief The apply function takes a vertex and a vector of numeric values
     (a total) as input. For unary vertices, this will be the sum of Lagrange
     multipliers, and we just need to sum that to the vertex potential and compute
     the argmax. For general factors, the vector of numeric values, as stated above,
     will contain all the Lagrange multipliers of the neighboring variables.
     So we need to loop through all possible factor configurations, get the
     sequence of states of each configuration, fetch the Lagrange multipliers for
     those states, and add them to the factor potential. Then we compute the argmax.
     It also computes dual and primal contribution for finding dual and primal
     objective values.
     */
    void apply(icontext_type& context, vertex_type& vertex, const gather_type& total)
    {   if (!opts.debug){
        vertex_data& vdata = vertex.data();

        if (vdata.nvars == 1)
        {
            // Unary factor.
            ASSERT_EQ(vdata.potentials.size(), total.messages.size());

            vec belief = vdata.potentials + total.messages;
             // Find primal contrib
            vdata.primal_contrib = vdata.potentials[vdata.best_configuration];
             // Save the best configuration for this vertex and find dual contrib
            vdata.dual_contrib = belief.maxCoeff(&vdata.best_configuration);

            if (opts.verbose > 1)
            {
                cout << "Vertex: " << vertex.id() << "\n";
                cout << "Potential: " << vdata.potentials << "\n";
                cout << "incomming message: " << total.messages << "\n";
                cout << "belief: " << belief << "\n";
                cout << "dual contrib: " << vdata.dual_contrib << "\n";
                cout << "vdata.best_configuration = " << vdata.best_configuration << "\n---\n";
            }
        }
        else
        {
            // General factor.
            vec belief = vdata.potentials;
            int num_configurations = vdata.potentials.size();
            for (int index_configuration = 0;
                 index_configuration < num_configurations;
                 ++index_configuration)
            {
                vector<int> states(vdata.nvars, -1);
                // This could be made more efficient by defining an iterator over factor
                // configurations.
                get_configuration_states(vertex, index_configuration, &states);
                int offset = 0;
                for (int k = 0; k < vdata.nvars; ++k)
                {
                    belief[index_configuration] += total.messages[offset + states[k]];
                    offset += vdata.cards[k];
                }
            }
            // Save the best configuration for this factor and find dual contrib
            vdata.dual_contrib = belief.maxCoeff(&vdata.best_configuration);
            //Find primal contrib
            int conf_index = get_configuration_index(vertex, total.neighbor_conf);
            vdata.primal_contrib = vdata.potentials[conf_index];
            //Find contribution fir sum of square of gradient
            vdata.sum_sq_norm_g = total.sq_norm_g;

            if (opts.verbose > 1)
            {
                cout << "Vertex: " << vertex.id() << "\n";
                cout << "Potential: " << vdata.potentials << "\n";
                cout << "incomming message: " << total.messages << "\n";
                cout << "belief: " << belief << "\n";
                cout << "dual contrib: " << vdata.dual_contrib << "\n";
                cout << "vdata.best_configuration = " << vdata.best_configuration << "\n---\n";
            }
        }
       }
      else usleep(1);

      if (opts.verbose > 1)
            cout << "end apply" << endl;
    }; // end of apply

    /**
     * \brief Since the MRF is undirected we will use all edges for gather and
     * scatter
     */
    edge_dir_type scatter_edges(icontext_type& context, const vertex_type& vertex) const
    {
         return graphlab::ALL_EDGES;
        // NOTENOTE: This assumes a sync engine.
    }; // end of scatter edges

    /**
     * \brief The scatter function takes a vertex and an edge as input.
     We just need to update the messages (Lagrange multipliers) by looking at the
     saved argmaxes.
     */
    void scatter(icontext_type& context, const vertex_type& vertex, edge_type& edge) const
    {
        const vertex_type other_vertex = get_other_vertex(edge, vertex);
        const vertex_type *unary_vertex;
        const vertex_type *factor_vertex;
     if (!opts.debug){
        if (vertex.data().nvars == 1)
        {
            // Unary factor.
            unary_vertex = &vertex;
            factor_vertex = &other_vertex;
        }
        else
        {
            // General factor.
            unary_vertex = &other_vertex;
            factor_vertex = &vertex;
        }
        const vertex_data& vdata = unary_vertex->data();
        const vertex_data& other_vdata = factor_vertex->data();
        edge_data& edata = edge.data();

        if (opts.verbose > 1)
            cout << "begin scatter" << endl;

        int iter_since_aggregate = (context.iteration()+2) - global_vars.iter_at_aggregate ;
        double stepsize = update_stepsize(context, 1, global_vars.old_dual, global_vars.primal_best,
                                      global_vars.sq_norm_g, global_vars.dual_inc_count, iter_since_aggregate);

        CHECK_GE(vdata.best_configuration, 0);
        CHECK_LT(vdata.best_configuration, vdata.cards[0]);

        // Negative subgradient

        edata.multiplier_messages[vdata.best_configuration] -= stepsize;

        vector<int> states(other_vdata.nvars, -1);
        get_configuration_states(*factor_vertex, other_vdata.best_configuration, &states);
        int offset = 0;
        int index_neighbor = -1;
        for (int k = 0; k < other_vdata.nvars; ++k) {
            int vertex_id = other_vdata.neighbors[k];
            if (vertex_id == unary_vertex->id()) {
                index_neighbor = k;
                break;
            }
            offset += other_vdata.cards[k];
        }
        CHECK_GE(index_neighbor, 0);
        CHECK_GE(states[index_neighbor], 0);
        CHECK_LT(states[index_neighbor], other_vdata.cards[index_neighbor]);
        CHECK_EQ(other_vdata.cards[index_neighbor], vdata.cards[0]);

        // Negative subgradient
        edata.multiplier_messages[states[index_neighbor]] += stepsize;

        //if (opts.verbose > 1)
        if (opts.verbose>1)
        {
            cout << "Scatter on (" << unary_vertex->id() << "," << factor_vertex->id() << ") called from " << vertex.id() << "\n";
            cout << "unary best config = " << vdata.best_configuration << "\n"
                 << "factor best config = " << states[index_neighbor] << "\n";
            cout << "Message: " << edata.multiplier_messages << "\n---\n";
        }
       }
        if (opts.verbose > 1)
            cout << "end scatter" << endl;

        // Signalling the other vertex and yourself to start.
        if ((context.iteration()+1) < opts.maxiter && global_vars.converged == false)
        {
            context.signal(vertex);
            context.signal(other_vertex);
        }

    }; // end of scatter

}; // end of class dd_vertex_program_symmetric

////////////////////////////////////////////////////////////////////////////////
// This class implements the "projected" version of dual decomposition described
// in:
// Komodakis, N., Paragios, N., and Tziritas, G. (2007).
// "MRF optimization via dual decomposition: Message-passing revisited"
// In Proc. of International Conference on Computer Vision.
//
// The formulation used is the one in Algorithm 1 of:
//
// André F. T. Martins, Mário A. T. Figueiredo, Pedro M. Q. Aguiar,
// Noah A. Smith, and Eric P. Xing.
// "An Augmented Lagrangian Approach to Constrained MAP Inference."
// International Conference on Machine Learning (ICML), 2011.
////////////////////////////////////////////////////////////////////////////////

struct dd_vertex_program_projected : public dd_vertex_program {

     graphlab::timer vertex_timer;
    /**
     * \brief Since the MRF is undirected we will use all edges for gather and
     * scatter
     */
    edge_dir_type gather_edges(icontext_type& context,
                               const vertex_type& vertex) const {
        if(!opts.debug){
        return graphlab::ALL_EDGES; }
        else
        return graphlab::NO_EDGES;
    }; // end of gather_edges

    /**
     * \brief The gather function takes a vertex and an edge as inputs and outputs
     a vector of numeric values. Vectors of numeric values will later be summed
     over all edges incident in this vertex.
     If the vertex is a unary factor, compute the sum of all the local MAP variables,
     which in the "apply" function will serve to compute the global MAP.
     Otherwise (if vertex is a general factor), things are a little more tricky.
     Suppose the factor is linked to K variables, with cardinalities C_1, ..., C_K.
     Suppose this edge is with respect to the k-th variable. Then, we return a
     vector of size C_1 + ... + C_K which is zero everywhere except in the
     k-th slot, where the Lagrange multipliers in "edge.messages" will be copied
     to. This way, when the "gather sum" takes place, and since all these slots
     are disjoint, we will just get the Lagrange multipliers of all the variables.
     */
    gather_type gather(icontext_type& context, const vertex_type& vertex,
                       edge_type& edge) const {
        const vertex_type other_vertex = get_other_vertex(edge, vertex);
        const vertex_data& vdata = vertex.data();
        edge_data& edata = edge.data();

        if (vdata.factor_type == VAR )
        {   // Unary factor.
            if (opts.verbose > 1) {
            cout << "This unary factor has " << vertex.num_in_edges() <<
            " in edges and " << vertex.num_out_edges() << " out edges" << endl; }
            gather_type gatherdata(edata.local_messages);
            return gatherdata;
        }
        else
        {
            // General factor.
            factor_type messages;
            //messages.resize(vdata.potentials.size());
            messages.setZero(vdata.potentials.size());
            int offset = 0;
            int index_neighbor = -1;
            for (int k = 0; k < vdata.nvars; ++k) {
                int vertex_id = vdata.neighbors[k];
                if (vertex_id == other_vertex.id()) {
                    index_neighbor = k;
                    break;
                }
                offset += vdata.cards[k];
            }

            CHECK_GE(index_neighbor, 0);
            vector <int> neighbor_conf(vdata.nvars, 0);
            neighbor_conf[index_neighbor] = other_vertex.data().best_configuration;
            for (int state = 0; state < vdata.cards[index_neighbor]; ++state) {
                messages[offset + state] = edata.multiplier_messages[state] + edata.potentials[state];
            }
            gather_type gather_data(messages,neighbor_conf);
            return gather_data;
        }
    }; // end of gather function

    /**
     * \brief The apply function takes a vertex and a vector of numeric values
     (a total) as input.
     For a unary vertex, "total" will be the sum of local MAP vectors, and we
     just need to divide by the vertex degree and save the result as global MAP.
     For higher-order factors, "total" will contain all the Lagrange multipliers
     of the neighboring variables. So we need to loop through all possible factor
     configurations, get the sequence of states of each configuration, fetch the
     Lagrange multipliers for those states, and add them to the factor potential.
     Then we compute the argmax and save result to local MAP for each variable
     connected to the factor.
     Note that since global MAP is computed from local MAP , locla MAP is needed
     to be computed before global MAP. Hence in even iterations local MAP is
     computed ( in scatter step) and in the subsequent iteration (which is hence
     odd) global MAP (in apply step) and multiplier messages (in scatter)
     are updated.
     */
    void apply(icontext_type& context, vertex_type& vertex,
               const gather_type& total) {
        vertex_data& vdata = vertex.data();

     if (!opts.debug){
        if (vdata.factor_type == VAR ) {
           if (context.iteration()%2 != 0) {
            vdata.beliefs = total.messages / static_cast<double>(vdata.degree);
            vdata.beliefs.maxCoeff(&vdata.best_configuration);
            //Find primal contrib
            vdata.primal_contrib = vdata.potentials[vdata.best_configuration];
            }
        }
        else {
          if(context.iteration()%2 == 0){
            // General factor.
            vec beliefs = vdata.potentials;
            int num_configurations = vdata.potentials.size();
            for (int index_configuration = 0;
                 index_configuration < num_configurations;
                 ++index_configuration) {
                vector<int> states(vdata.nvars, -1);
                // This could be made more efficient by defining an iterator over factor
                // configurations.
                get_configuration_states(vertex, index_configuration, &states);
                int offset = 0;
                for (int k = 0; k < vdata.nvars; ++k) {
                    beliefs[index_configuration] += total.messages[offset + states[k]];
                    offset += vdata.cards[k];
                }
            }
            // Save the best configuration for this factor and find dual contrib
            vdata.dual_contrib = beliefs.maxCoeff(&vdata.best_configuration);
            //Find primal contrib
            int conf_index = get_configuration_index(vertex, total.neighbor_conf);
            vdata.primal_contrib = vdata.potentials[conf_index];
         }
        }
      }
      else usleep(1);
    }; // end of apply

        /**
     * \brief Since the MRF is undirected we will use all edges for gather and
     * scatter
     */
    edge_dir_type scatter_edges(icontext_type& context,
                                const vertex_type& vertex) const {
        return graphlab::ALL_EDGES;
    }; // end of scatter edges

    /**
     * \brief The scatter function takes a vertex and an edge as input.
     (1) If the vertex is a unary factor, we update the messages (Lagrange multipliers)
     by subtracting the global MAP by the local MAP.
     (2) If the vertex is a higher order factor, this function will take the best
     configuration (obtained at the apply function) and save the local MAP
     at the corresponding edge.
     */
    void scatter(icontext_type& context, const vertex_type& vertex,
                 edge_type& edge) const {

        const vertex_data& vdata = vertex.data();
        edge_data& edata = edge.data();
        const vertex_type other_vertex = get_other_vertex(edge, vertex);
        if(!opts.debug){
        if (vdata.factor_type == VAR ) {
            if (context.iteration()%2 != 0) {

            // Unary factor. Update the messages (Lagrange multipliers).
             double stepsize = update_stepsize(context, 1, global_vars.old_dual, global_vars.primal_best,
                                              global_vars.sq_norm_g, global_vars.dual_inc_count, 0);

             edata.multiplier_messages += (vdata.beliefs - edata.local_messages) * stepsize;
          }
        }
       else
       {   if (context.iteration()%2 == 0) {
            //General factor. Update the local MAPs.
            const vertex_type &unary_vertex = get_other_vertex(edge, vertex);
            vector<int> states(vdata.nvars, -1);
            get_configuration_states(vertex, vdata.best_configuration, &states);
            int offset = 0;
            int index_neighbor = -1;
            for (int k = 0; k < vdata.nvars; ++k) {
                int vertex_id = vdata.neighbors[k];
                if (vertex_id == unary_vertex.id()) {
                    index_neighbor = k;
                    break;
                }
                offset += vdata.cards[k];
            }
            CHECK_GE(index_neighbor, 0);
            CHECK_GE(states[index_neighbor], 0);
            CHECK_LT(states[index_neighbor], vdata.cards[index_neighbor]);
            //CHECK_EQ(vdata.cards[index_neighbor], unary_vertex.data().cards[0]);
            edata.local_messages.setZero();
            edata.local_messages[states[index_neighbor]] += 1.0;
          }
        }
      }
        if ((context.iteration()+1) < opts.maxiter && global_vars.converged == false)
        {
            context.signal(vertex);
            context.signal(other_vertex);
        }
    }; // end of scatter
}; // end of class dd_vertex_program_projected

///////////////////////////////////////////////////////////////////////////////////
// This class implements the general Alternating Directions Method of Multipliers.
//
//  The formulation used is the one in Algorithm 2 of:
//
// André F. T. Martins, Mário A. T. Figueiredo, Pedro M. Q. Aguiar,
// Noah A. Smith, and Eric P. Xing.
// "Alternating Directions Dual Decomposition"
// Arxiv preprint arXiv:1212.6550, 2012.

///////////////////////////////////////////////////////////////////////////////////

struct admm_vertex_program:public dd_vertex_program {

  typedef int Configuration;

  // Function to solve each quadratic programming sub problem
  virtual void compute_beliefs(vertex_type& vertex,const gather_type& total,
                 vec& variable_posteriors, vec& additional_posteriors) = 0;
  virtual void SolveMAP(vertex_type& vertex,const gather_type& total,
                 vec& variable_posteriors, vec& additional_posteriors, double& value) = 0;

   /**
     * \brief Since the MRF is undirected we will use all edges for gather and
     * scatter
     */
    edge_dir_type gather_edges(icontext_type& context,
                               const vertex_type& vertex) const {
        if(!opts.debug){
        return graphlab::ALL_EDGES; }
        else
        return graphlab::NO_EDGES;
    }; // end of gather_edges

    /**
     * \brief The gather function takes a vertex and an edge as inputs and outputs
     a vector of numeric values. Vectors of numeric values will later be summed
     over all edges incident in this vertex.
     If the vertex is a unary factor, compute the sum of all the local MAP variables,
     which in the "apply" function will serve to compute the global MAP.
     Otherwise (if vertex is a general factor), things are a little more tricky.
     Suppose the factor is linked to K variables, with cardinalities C_1, ..., C_K.
     Suppose this edge is with respect to the k-th variable. Then, we return a
     vector of size C_1 + ... + C_K which is zero everywhere except in the
     k-th slot, where the Lagrange multipliers in "edge.messages" will be copied
     to. This way, when the "gather sum" takes place, and since all these slots
     are disjoint, we will just get the Lagrange multipliers of all the variables.
     */

   gather_type gather(icontext_type& context, const vertex_type& vertex,
                       edge_type& edge) const {
        const vertex_type other_vertex = get_other_vertex(edge, vertex);
        const vertex_data& vdata = vertex.data();
        edge_data& edata = edge.data();
        if (vdata.factor_type == VAR ) {
            // Unary factor.
            if (opts.verbose > 1){
            cout << "This unary factor has " << vertex.num_in_edges() <<
            " in edges and " << vertex.num_out_edges() << " out edges" << endl;
            }
            gather_type gatherdata(edata.local_messages);
            return gatherdata;
        }
        else if(vdata.factor_type != BUDGET){
            // Dense factor.
            factor_type messages, neighbor_distribution, multipliers;
            messages.setZero(vdata.potentials.size());
            neighbor_distribution.setZero(vdata.potentials.size());
            multipliers.setZero(vdata.potentials.size());
            int offset = 0;
            int index_neighbor = -1;
            for (int k = 0; k < vdata.nvars; ++k) {
                int vertex_id = vdata.neighbors[k];
                if (vertex_id == other_vertex.id()) {
                    index_neighbor = k;
                    break;}
                offset += vdata.cards[k];
            }
            CHECK_GE(index_neighbor, 0);

            vector <int> neighbor_conf(vdata.nvars, 0);
            neighbor_conf[index_neighbor] = other_vertex.data().best_configuration;

            for (int state = 0; state < vdata.cards[index_neighbor]; ++state) {
                messages[offset + state] = edata.multiplier_messages[state] + edata.potentials[state];
                multipliers[offset + state] = edata.multiplier_messages[state];
                neighbor_distribution[offset + state] = other_vertex.data().beliefs[state];
            }

            gather_type gather_data(messages,neighbor_conf);
            gather_data.neighbor_distribution = neighbor_distribution;
            gather_data.multipliers = multipliers;
            return gather_data;
        }

        else if(vdata.factor_type == BUDGET) {
            //Budget factor
            factor_type messages, neighbor_distribution;
            messages.setZero(vdata.nvars);
            neighbor_distribution.setZero(vdata.nvars);
            int index_neighbor = -1;
            for (int k = 0; k < vdata.nvars; ++k) {
                int vertex_id = vdata.neighbors[k];
                if (vertex_id == other_vertex.id()) {
                    index_neighbor = k;
                    break;}
            }
           messages[index_neighbor] =  edata.multiplier_messages[0] + edata.potentials[vdata.bound_states[index_neighbor]];

           neighbor_distribution[index_neighbor] = other_vertex.data().beliefs[vdata.bound_states[index_neighbor]];

           gather_type gather_data(messages);
           gather_data.neighbor_distribution = neighbor_distribution;
           return gather_data;

        }
    }; // end of gather function

   /**
     * \brief The apply function takes a vertex and a vector of numeric values
     (a total) as input.
     For a unary vertex, "total" will be the sum of local MAP vectors, and we
     just need to divide by the vertex degree and save the result as global MAP.
     For higher-order factors, "total" will contain all the Lagrange multipliers
     of the neighboring variables. So we need to loop through all possible factor
     configurations, get the sequence of states of each configuration, fetch the
     Lagrange multipliers for those states, and add them to the factor potential.
     Then we compute solution of Quadratic subproblem and save result to local
     MAP for each variable connected to the factor.
     Note that since global MAP is computed from local MAP , locla MAP is needed
     to be computed before global MAP. Hence in even iterations local MAP is
     computed ( in scatter step) and in the subsequent iteration (which is hence
     odd) global MAP (in apply step) and multiplier messages (in scatter)
     are updated.
     */

  void apply(icontext_type& context, vertex_type& vertex,
               const gather_type& total) {
        vertex_data& vdata = vertex.data();
     if (!opts.debug){
        if (vdata.factor_type == VAR ) {
             if (context.iteration()%2 != 0) {
            // Unary factor.
            //Find dual residual contrib
                vec dual_res_contrib;
                dual_res_contrib.setZero(vdata.cards[0]);
                for(int i=0; i<vdata.cards[0]; i++){

                   dual_res_contrib[i] =  (total.messages[i] / static_cast<double>(vdata.unary_degree[i]))
                                                                                - vdata.beliefs[i] ;
                   dual_res_contrib[i] = dual_res_contrib[i] * dual_res_contrib[i];
            // update global MAP
                   vdata.beliefs[i] = total.messages[i] / static_cast<double>(vdata.unary_degree[i]);
                }
                vdata.dual_res_contrib = dual_res_contrib.sum();
            // Find best configuration
                vdata.beliefs.maxCoeff(&vdata.best_configuration);
            //Find relaxed primal contribution
                vdata.primal_rel_contrib = vdata.potentials.dot(vdata.beliefs);
            // Find primal contribution
                vdata.primal_contrib = vdata.potentials[vdata.best_configuration];

            }
        }
        else{
            if(context.iteration()%2 == 0){
            // Dense and Budget factors
               vec additional_posteriors, variable_posteriors;
               additional_posteriors.setZero(vdata.potentials.size());
               variable_posteriors.setZero(vdata.potentials.size());
               if(vdata.factor_type == BUDGET){
                 additional_posteriors.setZero(vdata.nvars);
                 variable_posteriors.setZero(vdata.nvars);
                }
            //Find dual contrib
                SolveMAP(vertex, total, variable_posteriors, additional_posteriors, vdata.dual_contrib);
            // Find relaxed primal contribution
                if(vdata.factor_type == DENSE) {
                  vdata.primal_rel_contrib = vdata.potentials.dot(vdata.factor_beliefs);
                }
                else if(vdata.factor_type == XOR){
                  vdata.primal_rel_contrib = vdata.potentials.dot(vdata.beliefs);
                }
            //Find primal residual contribution
                vec primal_res_contrib = vdata.beliefs - total.neighbor_distribution;
                for(int i=0;i< vdata.beliefs.size(); i++){
                   primal_res_contrib[i] = primal_res_contrib[i]*primal_res_contrib[i];
                }
                vdata.primal_res_contrib = primal_res_contrib.sum();
            // Compute QP subproblem solution
                compute_beliefs(vertex, total, vdata.beliefs, vdata.factor_beliefs);
               //cout<<vertex.id()<<" "<<vdata.factor_type<<" "<<vdata.beliefs<<endl;
            //Find primal contrib
                if(vdata.factor_type == DENSE){
                int conf_index = get_configuration_index(vertex, total.neighbor_conf);
                vdata.primal_contrib = vdata.potentials[conf_index]; }

            }
          }
       }
      else usleep(1);
    }; // end of apply

       /**
     * \brief Since the MRF is undirected we will use all edges for gather and
     * scatter
     */
    edge_dir_type scatter_edges(icontext_type& context,
                                const vertex_type& vertex) const {
        return graphlab::ALL_EDGES;
    }; // end of scatter edges

   /**
     * \brief The scatter function takes a vertex and an edge as input.
     (1) If the vertex is a unary factor, we update the messages (Lagrange multipliers)
     by subtracting the global MAP by the local MAP.
     (2) If the vertex is a higher order factor, this function will take the best
     configuration (obtained at the apply function) and save the local MAP
     at the corresponding edge.
     */
    void scatter(icontext_type& context, const vertex_type& vertex,
                 edge_type& edge) const {
        const vertex_data& vdata = vertex.data();
        edge_data& edata = edge.data();
        const vertex_type other_vertex = get_other_vertex(edge, vertex);
        if (!opts.debug){
         // Unary factor. Update the messages (Lagrange multipliers).
            if (vdata.factor_type == VAR ){
                if (context.iteration()%2 != 0) {

                    double stepsize = update_stepsize(context, 0 , global_vars.old_dual, global_vars.primal_best,
                                                       global_vars.sq_norm_g, global_vars.dual_inc_count, 0);
                    if(other_vertex.data().factor_type == BUDGET){
                       int index_neighbor = -1;
                       for(int i=0; i < other_vertex.data().nvars; i++) {
                           if(other_vertex.data().neighbors[i] == vertex.id()){
                              index_neighbor =i;
                              break;
                           }
                       }
                      edata.multiplier_messages[0] +=(vdata.beliefs[other_vertex.data().bound_states[index_neighbor]]
                         - edata.local_messages[other_vertex.data().bound_states[index_neighbor]]) * stepsize;
                     }
                     else if(other_vertex.data().factor_type != BUDGET){
                            edata.multiplier_messages += (vdata.beliefs - edata.local_messages) * stepsize;
                     }
                }
             }
             else if(vdata.factor_type != BUDGET){
                     if (context.iteration()%2 == 0) {
            //General factor. Update the local MAPs.
                        const vertex_type &unary_vertex = get_other_vertex(edge, vertex);
                        int offset = 0;
                        int index_neighbor = -1;
                        for (int k = 0; k < vdata.nvars; ++k) {
                            int vertex_id = vdata.neighbors[k];
                            if (vertex_id == unary_vertex.id()) {
                                index_neighbor = k;
                                break;
                            }
                            offset += vdata.cards[k];
                        }
                        CHECK_GE(index_neighbor, 0);
                        //CHECK_EQ(vdata.cards[index_neighbor], unary_vertex.data().cards[0]);

                        for(int state = 0; state < vdata.cards[index_neighbor]; state++){
                           edata.local_messages[state] = vdata.beliefs[offset+state];
                        }
                   }
              }
              else if(vdata.factor_type == BUDGET){
               // Budget factor. Update local MAPs
                      if (context.iteration()%2 == 0) {
                         const vertex_type &unary_vertex = get_other_vertex(edge, vertex);
                         edata.local_messages.setZero();
                         int index_neighbor = -1;
                         for (int k = 0; k < vdata.nvars; ++k) {
                             int vertex_id = vdata.neighbors[k];
                             if (vertex_id == unary_vertex.id()) {
                                index_neighbor = k;
                                break;
                             }
                         }
                         edata.local_messages[vdata.bound_states[index_neighbor]] =
                                                    vdata.beliefs[index_neighbor];
                      }
               }
        }
        if ((context.iteration()+1) < opts.maxiter && global_vars.converged == false) {
            context.signal(vertex);
            context.signal(other_vertex);
        }
    }; // end of scatter

 }; /* end of admm_vertex_program */

/////////////////////////////////////////////////////////////////////////////////
//Aggregator functions to compute primal & dual objectives and residuals

objective sum(dd_vertex_program::icontext_type& context, const dd_vertex_program::vertex_type& vertex){
  objective retval;
  retval.primal = vertex.data().primal_contrib;
  retval.dual = vertex.data().dual_contrib;
  retval.sum_sq_norm_g = vertex.data().sum_sq_norm_g;
  retval.primal_rel = vertex.data().primal_rel_contrib;
  retval.dual_res  = vertex.data().dual_res_contrib;
  retval.primal_res = vertex.data().primal_res_contrib;
  retval.total_confs =(vertex.data().nvars ==1)?(std::accumulate(vertex.data().unary_degree.begin(),
                                                  vertex.data().unary_degree.end(),0)):0;
  global_vars.iter_at_aggregate = (context.iteration() +2);

  return retval;
}

void print_obj(dd_vertex_program::icontext_type& context, objective total) {   //TODO make this function part of dd_vertex_program class
      if (context.iteration() % 2 == 0 || opts.algorithm == 0) {
        //saving best values
         if (total.dual > global_vars.old_dual){
             global_vars.dual_inc_count ++;
         }

         if (total.primal> global_vars.primal_best) {
            global_vars.primal_best = total.primal;
         }

         if(opts.verbose >0) {
            cout<<"iteration: "<<context.iteration()<<" Dual Objective: " << total.dual<<
                   " "<<"Primal Objective: "<<total.primal<<endl;
            cout<< "Best Primal so far: "  << global_vars.primal_best<<" "<<endl;
        //residuals only available in case of ad3
            if(opts.algorithm == 2) {
              cout<<"Relaxed Primal Objective:"<<total.primal_rel<<endl;
              cout<<"Dual Residual:"<< sqrt(total.dual_res/total.total_confs)<<" "
              <<"Primal Residual:"<<sqrt(total.primal_res/total.total_confs)<<endl;
            }
        }
        //convergence criteria for dd, bethe admm
        if ((std::fabs(total.dual-global_vars.old_dual) < opts.dualimprovthres) && opts.algorithm != 2) {
            global_vars.converged = true;
           cout<< "Dual Objective: " << total.dual<< " "<<"Primal Objective: "<<total.primal<<endl;
           cout<<" Number of iteration at convergence:"<<context.iteration() +2 <<endl;
        }
        //convergence criteria for ad3
        if((sqrt(total.dual_res/total.total_confs) < opts.dualimprovthres
             && sqrt(total.primal_res/total.total_confs) < opts.dualimprovthres)
      && (std::fabs(total.dual-global_vars.old_dual) < opts.dualimprovthres) && opts.algorithm == 2){
           global_vars.converged = true;
           cout<< "Dual Objective: " << total.dual<< " "<<"Primal Objective: "<<total.primal<<endl;
           cout<<" Number of iteration at convergence:"<<context.iteration() +2 <<endl;
        }

        global_vars.old_dual = total.dual;
        global_vars.sq_norm_g = total.sum_sq_norm_g;
       //saving history
        if (opts.history_file != "\0"){
            global_vars.history[0].push_back(context.iteration()+2);
            global_vars.history[1].push_back(global_vars.timer.current_time());
            global_vars.history[2].push_back(total.dual);
            global_vars.history[3].push_back(total.primal);
        }
    }

}

/* end of aggregator functions */

#endif


/* CirKit: A circuit toolkit
 * Copyright (C) 2009-2015  University of Bremen
 * Copyright (C) 2015-2017  EPFL
 *
 * Permission is hereby granted, free of charge, to any person
 * obtaining a copy of this software and associated documentation
 * files (the "Software"), to deal in the Software without
 * restriction, including without limitation the rights to use,
 * copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the
 * Software is furnished to do so, subject to the following
 * conditions:
 *
 * The above copyright notice and this permission notice shall be
 * included in all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
 * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
 * OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
 * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
 * HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
 * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
 * OTHER DEALINGS IN THE SOFTWARE.
 */

#include "bdd.hpp"

#include <iostream>
#include <vector>

#include <boost/lexical_cast.hpp>

#include <cuddObj.hh>

#include <alice/rules.hpp>

#include <cli/stores.hpp>
#include <core/utils/range_utils.hpp>
#include <core/utils/string_utils.hpp>

using namespace boost::program_options;

namespace cirkit
{

bdd_command::bdd_command( const environment::ptr& env )
  : cirkit_command( env, "BDD manipulation" )
{
  opts.add_options()
    ( "characteristic,c", value( &characteristic ), "Compute characteristic function (x: inputs first, y: outputs first)" )
    ( "clique",           value( &clique ),         "Computes clique(n,k) function, give n,k as string" )
    ( "new,n",                                      "Add a new entry to the store; if not set, the current entry is overriden" )
    ;
}

command::rules_t bdd_command::validity_rules() const
{
  return { has_store_element_if_set<bdd_function_t>( *this, env, "characteristic" ) };
}

bool bdd_command::execute()
{
  auto& bdds = env->store<bdd_function_t>();

  if ( is_set( "characteristic" ) )
  {
    auto bdd = bdds.current();

    extend_if_new( bdds );

    bdds.current() = compute_characteristic( bdd, characteristic == "x" );
  }

  else if ( is_set( "clique" ) )
  {
    const auto nk = split_string_pair( clique, "," );
    const auto n = boost::lexical_cast<unsigned>( nk.first );
    const auto k = boost::lexical_cast<unsigned>( nk.second );

    Cudd mgr;
    auto func = mgr.bddZero();

    lexicographic_combinations( n, k, [&mgr, &func, k]( const std::vector<unsigned>& v ) {
        auto term = mgr.bddOne();
        lexicographic_combinations( k, 2u, [&mgr, &term, &v]( const std::vector<unsigned>& idxs ) {
            const auto v0 = v[idxs[0u]];
            const auto v1 = v[idxs[1u]];
            const auto edge_idx = v0 + ( v1 * ( v1 - 1 ) ) / 2;

            term &= mgr.bddVar( edge_idx );

            return false;
          } );

        func |= term;

        return false;
      } );

    extend_if_new( bdds );
    bdds.current() = {mgr, {func}};
  }

  return true;
}

}

// Local Variables:
// c-basic-offset: 2
// eval: (c-set-offset 'substatement-open 0)
// eval: (c-set-offset 'innamespace 0)
// End:


#include <boost/foreach.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <iomanip>

#include "HoughTransform.hpp"
#include "mlp_threshold.hpp"
#include "hough_histogram.hpp"
#include "ImageProcessor.hpp"
#include "utils.hpp"
#include "draw_hist.hpp"
#include "detection.hpp"

#define foreach BOOST_FOREACH

namespace artelab
{

    ImageProcessor::ImageProcessor(std::string mlp_file, cv::Size win_size, std::string outdir, bool quiet, bool show) :
        _show(show),
        _quiet(quiet),
        _winsize(win_size)
    {
        _mlp.load(mlp_file);
        _output = DirectoryInfo(outdir);
    }

    ImageProcessor::~ImageProcessor()
    { }

    ImageProcessor& ImageProcessor::show(bool b)
    {
        _show = b;
        return *this;
    }

    ImageProcessor& ImageProcessor::set_output(std::string outdir)
    {
        _output = DirectoryInfo(outdir);
        return *this;
    }

    void ImageProcessor::show_image(std::string name, cv::Mat img)
    {
        if(_show)
        {
            cv::namedWindow(name);
            cv::imshow(name, img);
        }
    }

    void draw_lines_at_angle(double angle, std::vector<cv::Vec4i> lines, cv::Mat& image, int tolerance=2)
    {
        cv::Scalar color = image.channels() == 1? cv::Scalar(255) : cv::Scalar(0,0,255);

        foreach(cv::Vec4i l, lines)
        {
            double delta = double(l[3] - l[1]) / double(l[2] - l[0]);
            double act_angle = atan(delta)*180/CV_PI;
            act_angle += act_angle < 0? 180 : 0;
            double diff = fabs(act_angle - angle);
            if(diff < tolerance)
            {
                int thick = int(diff) > 3? 1 : 4-int(diff);
                cv::line(image, cv::Point(l[0], l[1]), cv::Point(l[2], l[3]), color, thick, 8);
            }
        }
    }

    void smooth_histogram(cv::Mat& hist, int kernel_size, int iterations)
    {
        CV_Assert(hist.cols == 1);

        int hist_type = hist.type();
        bool cvt_back = hist_type != CV_32F;
        if(cvt_back) hist.convertTo(hist, CV_32F);

        cv::Mat kernel = cv::Mat::ones(kernel_size, 1, CV_32F);
        kernel = kernel / float(kernel_size);

        for(int i=0; i < iterations; i++)
        {
            cv::filter2D(hist, hist, -1, kernel);
        }

        if(cvt_back) hist.convertTo(hist, hist_type);
    }

    void histograms_from_hough_lines(cv::Mat feature, cv::Mat& row_hist, cv::Mat& col_hist, bool smooth_and_thresh=true)
    {
        row_hist = get_histogram(feature, HIST_ROW, CV_32F);
        col_hist = get_histogram(feature, HIST_COL, CV_32F);

        if(smooth_and_thresh)
        {
            smooth_histogram(row_hist, 5, 200);
            smooth_histogram(col_hist, 5, 200);

            double row_hist_mean = cv::sum(row_hist)[0] / double(row_hist.rows);
            double col_hist_mean = cv::sum(col_hist)[0] / double(col_hist.rows);
            cv::threshold(row_hist, row_hist, row_hist_mean, 0, cv::THRESH_TOZERO);
            cv::threshold(col_hist, col_hist, col_hist_mean, 0, cv::THRESH_TOZERO);
        }
    }

    cv::Mat project_histograms(cv::Mat row_hist, cv::Mat col_hist)
    {
        CV_Assert(row_hist.type() == CV_32F);
        CV_Assert(col_hist.type() == CV_32F);

        cv::Mat out = row_hist * col_hist.t();

        double min, max;
        cv::minMaxLoc(out, &min, &max);
        out = (out - min) / max * 255;
        out.convertTo(out, CV_8U);

        return out;
    }

    results ImageProcessor::process(std::string name, ArtelabDataset::barcode_image bcimage)
    {
        cv::Mat img_orig = cv::imread(bcimage.original.fullName(), CV_LOAD_IMAGE_COLOR);
        cv::Mat img_truth = cv::imread(bcimage.detection_gt.fullName(), CV_LOAD_IMAGE_GRAYSCALE);

        TimeCounter tc;
        tc.start();

        // Apply canny
        cv::Mat img_canny;
        cv::GaussianBlur(img_orig, img_canny, cv::Size(17,17), 2);
        cv::Canny(img_canny, img_canny, 60, 100, 3);

        // Get hough transform
        HoughTransform hough(img_canny);

        // Threshold with MLP
        cv::Mat img_neural = threshold_mlp(_mlp, _winsize, hough);

        // find angle
        double angle = max_angle_hist(get_histogram(img_neural, HIST_ROW));

        // lines from canny using probabilistc hough
        const int tolerance = 3;
        std::vector<cv::Vec4i> lines;
        cv::HoughLinesP(img_canny, lines, 1, CV_PI/180, 50, 20, 1);
        cv::Mat line_image;
        if(_show)
        {
            // impress lines of given angle on the canny image
            cv::cvtColor(img_canny, line_image, CV_GRAY2BGR);
            draw_lines_at_angle((int(angle+0.5) + 90) % 180, lines, line_image, tolerance);
        }

        // obtain a feature image with only the lines. It is rectified.
        cv::Mat feature_image;
        feature_image = cv::Mat::zeros(img_canny.size(), CV_8U);
        draw_lines_at_angle((int(angle+0.5) + 90) % 180, lines, feature_image, tolerance);
        feature_image = rotate_image(feature_image, angle);

        //Histograms for detection
        cv::Mat row_hist, col_hist,feature_with_hist_smooth;
        histograms_from_hough_lines(feature_image, row_hist, col_hist);
        feature_with_hist_smooth = draw_histogram_on_image(row_hist, feature_image, cv::Scalar(0,0,255), HIST_ROW);
        feature_with_hist_smooth = draw_histogram_on_image(col_hist, feature_with_hist_smooth, cv::Scalar(0,255,0), HIST_COL);

        // project histograms
        cv::Mat img_hist_projection = project_histograms(row_hist, col_hist);

        // crop image with projected mask and threshold
        cv::Mat bb_mask;
        bb_mask = artelab::rotate_image(img_hist_projection, -angle);
        double min, max;
        cv::minMaxLoc(bb_mask, &min, &max);
        cv::threshold(bb_mask, bb_mask, int(0.3*max), 1, cv::THRESH_BINARY);
        cv::Mat img_cropped;
        if(_show) img_orig.copyTo(img_cropped, bb_mask);

        // draw bounding boxes
        std::vector<cv::Rect> rects = object_rectangles(img_hist_projection, int(0.3*max));
        cv::Mat img_bb;
        img_orig.copyTo(img_bb);
        cv::Mat img_detection_mask = cv::Mat::zeros(img_orig.size(), CV_8U);
        foreach(cv::Rect r, rects)
        {
            // BB for impression on original image
            cv::Mat bb = cv::Mat::zeros(img_bb.size(), CV_8U);
            cv::rectangle(bb, r, cv::Scalar(255), 2);
            bb = artelab::rotate_image(bb, -angle);
            img_bb.setTo(cv::Scalar(0,0,255), bb);

            // BB for measuring accuracy
            cv::Mat bb_fill = cv::Mat::zeros(img_orig.size(), CV_8U);
            cv::rectangle(bb_fill, r, cv::Scalar(255), -1);
            bb_fill = artelab::rotate_image(bb_fill, -angle);
            img_detection_mask.setTo(cv::Scalar(255), bb_fill);
        }

        tc.stop();

        // Measuring accuracy
        results res = measure_results(img_detection_mask, img_truth, tc);

        // verbose output
        if(!_quiet)
        {
            std::cout << std::setw(15) << std::left << name;
            std::cout << std::setw(11) << std::left << ("Angle " + tostring(angle));
            std::cout << std::setw(20) << std::left << ("Accuracy: " + tostring(res.jaccard));
            std::cout << std::setw(6) << std::left << ("Time: " + tostring(res.time)) << std::endl;
        }

        // show results
        show_image("Prob Hough", line_image);
        show_image("histograms smooth", feature_with_hist_smooth);
        show_image("histograms projection", img_hist_projection);
        show_image("Cropped", img_cropped);
        show_image("Boundig Boxes", img_bb);
        show_image("Original", img_orig);
        show_image("Feature", feature_image);
        show_image("Canny", img_canny);

        // saving intermediate images
        if(_output.fullPath() != "")
        {
            std::ostringstream ss;

            // Original
            ss << name << "_1_original.png";
            cv::imwrite(_output.fileCombine(ss.str()).fullName(), img_orig);

            // Hough accumulator
            ss.str(""); ss.clear(); ss << name << "_2_accumulator.png";
            cv::imwrite(_output.fileCombine(ss.str()).fullName(), hough.get_hough_image());

            // Hough MLP threshold
            img_neural = draw_histogram_on_image(get_histogram(img_neural, HIST_ROW, CV_32F), img_neural, cv::Scalar(0,0,255), HIST_ROW);
            ss.str(""); ss.clear();; ss << name << "_3_thresh_mlp.png";
            cv::imwrite(_output.fileCombine(ss.str()).fullName(), img_neural);

            // Feature with histograms not processed
            histograms_from_hough_lines(feature_image, row_hist, col_hist, false);
            cv::Mat feature_with_hist = draw_histogram_on_image(row_hist, feature_image, cv::Scalar(0,0,255), HIST_ROW);
            feature_with_hist = draw_histogram_on_image(col_hist, feature_with_hist, cv::Scalar(0,255,0), HIST_COL);
            ss.str(""); ss.clear();; ss << name << "_4_hist.png";
            cv::imwrite(_output.fileCombine(ss.str()).fullName(), feature_with_hist);

            // Feature with processed histograms
            ss.str(""); ss.clear();; ss << name << "_5_hist_smooth_thresh.png";
            cv::imwrite(_output.fileCombine(ss.str()).fullName(), feature_with_hist_smooth);

            // Histogram projection
            ss.str(""); ss.clear();; ss << name << "_6_hist_projection.png";
            cv::imwrite(_output.fileCombine(ss.str()).fullName(), img_hist_projection);

            // Bounding boxes
            ss.str(""); ss.clear();; ss << name << "_7_boxes.png";
            cv::imwrite(_output.fileCombine(ss.str()).fullName(), img_bb);

            // Canny
            ss.str(""); ss.clear(); ss << name << "_canny.png";
            cv::imwrite(_output.fileCombine(ss.str()).fullName(), img_canny);
        }

        return res;
    }

}


//=======================================================================
// Copyright (c) 2013 Maciej Andrejczuk
//               2014 Piotr Wygocki
//
// 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)
//=======================================================================
/**
 * @file steiner_utils.hpp
 * @brief
 * @author Maciej Andrejczuk, Piotr Wygocki
 * @version 1.0
 * @date 2013-08-01
 */
#ifndef PAAL_STEINER_UTILS_HPP
#define PAAL_STEINER_UTILS_HPP

#define BOOST_RESULT_OF_USE_DECLTYPE

#include "paal/data_structures/bimap.hpp"
#include "paal/data_structures/metric/metric_to_bgl.hpp"
#include "paal/utils/irange.hpp"

#include <boost/graph/prim_minimum_spanning_tree.hpp>
#include <boost/range/join.hpp>

namespace paal {
namespace ir {

/**
 * Utilities class. Responsible for calculating the cost of given edges.
 */
class steiner_utils {
  public:
    /**
     * Calculates total cost of given edges.
     */
    template <typename Metric, typename Terminals, typename Result>
    static typename data_structures::metric_traits<Metric>::DistanceType
    count_cost(const Result& steiner_vertices, const Terminals& terminals, const Metric& cost_map) {
        using Vertex = typename data_structures::metric_traits<Metric>::VertexType;
        using Dist = typename data_structures::metric_traits<Metric>::DistanceType;

        auto all_elements = boost::range::join(terminals, steiner_vertices);
        paal::data_structures::bimap<Vertex> idx;
        auto g = paal::data_structures::metric_to_bgl_with_index(cost_map,
                all_elements, idx);
        std::vector<std::size_t> pm(all_elements.size());
        boost::prim_minimum_spanning_tree(g, &pm[0]);
        auto idx_m = paal::data_structures::make_metric_on_idx(cost_map, idx);

        Dist cost{};
        for (auto i : irange(pm.size())) {
            cost += idx_m(i, pm[i]);
        }
        return cost;
    }
};

} // ir
} // paal

#endif // PAAL_STEINER_UTILS_HPP


#ifndef NEPOMUK_TOOL_CONTAINER_SURJECTIVE_MAPPING_HPP_
#define NEPOMUK_TOOL_CONTAINER_SURJECTIVE_MAPPING_HPP_

#include <boost/range/iterator_range.hpp>
#include <cstddef>
#include <vector>

namespace nepomuk
{
namespace tool
{
namespace container
{

// a surjective mapping is a mapping that applies an image to all items in a codomain x into a
// domain y.
template <typename codomain_type, typename domain_type> class SurjectiveMapping
{
  public:
    using iterator = typename std::vector<codomain_type>::const_iterator;
    using iterator_range = typename boost::iterator_range<iterator>;

    SurjectiveMapping(std::size_t const codomain_size, std::size_t domain_size);

    // return the image of a codomain type
    domain_type codomain_mapping(codomain_type const value) const;

    // return the preimage range of a domain type
    iterator_range domain_mapping(domain_type const value) const;

    // add stuff to the mapping
    void map_codomain(codomain_type const from, domain_type const to);

    std::size_t codomain_size() const;
    std::size_t domain_size() const;

  private:
    // reverse mapping
    std::vector<std::vector<codomain_type>> domain_mappings;
    // forward mapping (surjectivity guarantees a single entry per domain type here)
    std::vector<domain_type> codomain_mappings;
};

template <typename codomain_type, typename domain_type>
domain_type
SurjectiveMapping<codomain_type, domain_type>::codomain_mapping(codomain_type const value) const
{
    return codomain_mappings[static_cast<std::uint64_t>(value)];
}

template <typename codomain_type, typename domain_type>
typename SurjectiveMapping<codomain_type, domain_type>::iterator_range
SurjectiveMapping<codomain_type, domain_type>::domain_mapping(domain_type const value) const
{
    return boost::make_iterator_range(domain_mappings[static_cast<std::uint64_t>(value)].begin(),
                                      domain_mappings[static_cast<std::uint64_t>(value)].end());
}

template <typename codomain_type, typename domain_type>
SurjectiveMapping<codomain_type, domain_type>::SurjectiveMapping(std::size_t const codomain_size,
                                                                 std::size_t domain_size)
    : domain_mappings(domain_size), codomain_mappings(codomain_size)
{
    // resizes the mappings
}

template <typename codomain_type, typename domain_type>
void SurjectiveMapping<codomain_type, domain_type>::map_codomain(codomain_type const from,
                                                                 domain_type const to)
{
    // add a new mapping to the domain
    codomain_mappings[static_cast<std::uint64_t>(from)] = to;
    domain_mappings[static_cast<std::uint64_t>(to)].push_back(from);
}

template <typename codomain_type, typename domain_type>
std::size_t SurjectiveMapping<codomain_type, domain_type>::domain_size() const
{
    return domain_mappings.size();
}

template <typename codomain_type, typename domain_type>
std::size_t SurjectiveMapping<codomain_type, domain_type>::codomain_size() const
{
    return codomain_mappings.size();
}

} // namespace container
} // namespace tool
} // namespace nepomuk

#endif // NEPOMUK_TOOL_CONTAINER_SURJECTIVE_MAPPING_HPP_


#pragma once

#include <Eigen/Core>
#include <Eigen/Sparse>

#include <ipc/collision_constraint.hpp>
#include <ipc/friction/friction_constraint.hpp>
#include <ipc/utils/eigen_ext.hpp>

namespace ipc {

void construct_friction_constraint_set(
    const Eigen::MatrixXd& V,
    const Eigen::MatrixXi& E,
    const Eigen::MatrixXi& F,
    const Constraints& contact_constraint_set,
    double dhat,
    double barrier_stiffness,
    double mu,
    FrictionConstraints& friction_constraint_set);

void construct_friction_constraint_set(
    const Eigen::MatrixXd& V,
    const Eigen::MatrixXi& E,
    const Eigen::MatrixXi& F,
    const Constraints& contact_constraint_set,
    double dhat,
    double barrier_stiffness,
    const Eigen::VectorXd& mus,
    FrictionConstraints& friction_constraint_set);

void construct_friction_constraint_set(
    const Eigen::MatrixXd& V,
    const Eigen::MatrixXi& E,
    const Eigen::MatrixXi& F,
    const Constraints& contact_constraint_set,
    double dhat,
    double barrier_stiffness,
    const Eigen::VectorXd& mus,
    const std::function<double(double, double)>& blend_mu,
    FrictionConstraints& friction_constraint_set);

/// @brief Compute the friction potential between to positions.
///
/// @param V0 Vertex positions at start of time-step (rowwise)
/// @param V1 Current vertex positions (rowwise)
/// @param E  Edge vertex indicies
/// @param F  Face vertex indicies (empty in 2D)
/// @param friction_constraint_set
/// @param epsv_times_h
template <typename T>
T compute_friction_potential(
    const Eigen::MatrixXd& V0,
    const Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>& V1,
    const Eigen::MatrixXi& E,
    const Eigen::MatrixXi& F,
    const FrictionConstraints& friction_constraint_set,
    double epsv_times_h);

Eigen::VectorXd compute_friction_potential_gradient(
    const Eigen::MatrixXd& V0,
    const Eigen::MatrixXd& V1,
    const Eigen::MatrixXi& E,
    const Eigen::MatrixXi& F,
    const FrictionConstraints& friction_constraint_set,
    double epsv_times_h);

Eigen::SparseMatrix<double> compute_friction_potential_hessian(
    const Eigen::MatrixXd& V0,
    const Eigen::MatrixXd& V1,
    const Eigen::MatrixXi& E,
    const Eigen::MatrixXi& F,
    const FrictionConstraints& friction_constraint_set,
    double epsv_times_h,
    bool project_hessian_to_psd = true);

} // namespace ipc

#include "friction.tpp"


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


/**
 * @file penalty_impl.hpp
 * @author Leonardo Arcari (leonardo1.arcari@gmail.com)
 * @version 1.0.0
 * @date 2018-10-28
 *
 * @copyright Copyright (c) 2018 Leonardo Arcari
 *
 * MIT License
 *
 * 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 BOOST_PENALTY_BASED_ALTERNATIVE_ROUTING_IMPL_HPP
#define BOOST_PENALTY_BASED_ALTERNATIVE_ROUTING_IMPL_HPP

#include <boost/graph/dijkstra_shortest_paths.hpp>
#include <boost/graph/graph_concepts.hpp>
#include <boost/graph/graph_traits.hpp>
#include <boost/graph/properties.hpp>
#include <boost/graph/reverse_graph.hpp>
#include <boost/property_map/function_property_map.hpp>
#include <boost/property_map/property_map.hpp>

#include <arlib/details/arlib_utils.hpp>
#include <arlib/routing_kernels/bidirectional_dijkstra.hpp>
#include <arlib/routing_kernels/types.hpp>
#include <arlib/terminators.hpp>
#include <arlib/type_traits.hpp>

#include <functional>
#include <iostream>
#include <queue>
#include <unordered_map>
#include <unordered_set>
#include <vector>

namespace arlib {
/**
 * Implementations details of kSPwLO algorithms
 */
namespace details {
//===----------------------------------------------------------------------===//
//                      Penalty algorithm types
//===----------------------------------------------------------------------===//
/**
 * A map from edges to the number of time they have been penalized so
 * far.
 *
 * @tparam Edge An edge_descriptor
 */
template <typename Edge>
using PenBoundsMap = std::unordered_map<Edge, int, boost::hash<Edge>>;

/**
 * A map from edges to their weights
 *
 * @tparam Edge An edge_descriptor
 * @tparam Length The edge weight type
 */
template <typename Edge, typename Length>
using WeightMap = std::unordered_map<Edge, Length, boost::hash<Edge>>;

/**
 * A vector for tracking distance of a vertex from the source.
 *
 * The vector must be indexable by the vertex_descriptor, thus its size should
 * be equal to the number of vertices in the graph.
 *
 * @tparam Length The edge weight type.
 */
template <typename Length> using DistanceMap = std::vector<Length>;

//===----------------------------------------------------------------------===//
//                      Penalty algorithm classes
//===----------------------------------------------------------------------===//

/**
 * A functor to return the penalized weight of an edge, to avoid changing
 * the original graph weights.
 *
 * @tparam PMap A Weight Property Map.
 */
template <typename PMap> class penalty_functor {
public:
  using Edge = typename boost::property_traits<PMap>::key_type;
  using Length = double;

  /**
   * Construct a new penalty functor object.
   *
   * @tparam EdgeIterator An iterator of the graph edges.
   * @param weight The weight property map.
   */
  penalty_functor(PMap weight)
      : weight{weight}, penalties{std::make_shared<WeightMap<Edge, Length>>()} {
  }

  /**
   * Copy constructor.
   *
   * @param other The penalty functor to copy from
   */
  penalty_functor(const penalty_functor<PMap> &other)
      : weight{other.weight}, penalties{other.penalties} {}

  /**
   * Returns the penalized weight of an edge.
   *
   * @param e The query edge.
   * @return The penalized weight for @p e.
   */
  const Length &operator()(const Edge &e) const { return get_or_insert(e); }

  /**
   * Returns the penalized weight of an edge.
   *
   * @param e The query edge.
   * @return The penalized weight for @p e.
   */
  Length &operator[](const Edge &e) { return get_or_insert(e); }

  /**
   * Returns the penalized weight of an edge.
   *
   * @param e The query edge.
   * @return The penalized weight for @p e.
   */
  const Length &operator[](const Edge &e) const { return get_or_insert(e); }

  penalty_functor clone() const {
    auto pf = *this;
    pf.penalties = std::make_shared<WeightMap<Edge, Length>>(*penalties);
    return pf;
  }

private:
  PMap weight;
  mutable std::shared_ptr<WeightMap<Edge, Length>> penalties;

  Length &get_or_insert(const Edge &e) const {
    if (auto search = penalties->find(e); search != penalties->end()) {
      return search->second;
    } else {
      auto [it, ok] = penalties->insert({e, weight[e]});
      assert(ok && "[kspwlo::penalty_functor] Could not insert edge weight");
      return it->second;
    }
  }
};

template <typename PMap, typename Graph> class reverse_penalty_functor {
public:
  using Edge = typename boost::graph_traits<
      boost::reverse_graph<Graph>>::edge_descriptor;
  using Length = double;

  reverse_penalty_functor(penalty_functor<PMap> &penalty, const Graph &G,
                          const boost::reverse_graph<Graph> &rev_G)
      : inner_pf{penalty}, G{G}, rev_G{rev_G} {}

  reverse_penalty_functor(reverse_penalty_functor const &other)
      : inner_pf{other.inner_pf}, G{other.G}, rev_G{other.rev_G} {}

  const Length &operator()(const Edge &e) const {
    auto forward_edge = get_forward_edge(e);

    return inner_pf(forward_edge);
  }

  Length &operator[](const Edge &e) {
    auto forward_edge = get_forward_edge(e);

    return inner_pf[forward_edge];
  }

  const Length &operator[](const Edge &e) const {
    auto forward_edge = get_forward_edge(e);

    return inner_pf[forward_edge];
  }

private:
  auto get_forward_edge(const Edge &e) const {
    using namespace boost;
    auto u = source(e, rev_G);
    auto v = target(e, rev_G);

    auto [forward_edge, is_valid] = edge(v, u, G);
    assert(is_valid);
    return forward_edge;
  }

  penalty_functor<PMap> &inner_pf;
  const Graph &G;
  const boost::reverse_graph<Graph> &rev_G;
};

//===----------------------------------------------------------------------===//
//                     Penalty algorithm support routines
//===----------------------------------------------------------------------===//
/**
 * Computes the shortest path between two vertices s and t, first
 * from s to t and then from t to s. The distances of each node in the
 * shortest paths are stored in distance_s and distance_t.
 *
 * @tparam Graph A Boost::PropertyGraph having at least one edge
 *               property with tag boost::edge_weight_t.
 * @tparam Vertex A vertex_descriptor
 * @tparam Length The edge weight type.
 * @param G The graph.
 * @param s The source vertex.
 * @param t The target vertex.
 * @param distance_s A map from Vertex to its distance from @p s
 * @param distance_t A map from Vertex to its distance from @p t
 * @return A vector of the edges of the shortest path from s to t.
 *         An empty optional if t is not reachable from s.
 */
template <typename Graph, typename Vertex = vertex_of_t<Graph>,
          typename Edge = edge_of_t<Graph>,
          typename Length = length_of_t<Graph>>
std::optional<std::vector<Edge>>
dijkstra_shortest_path_two_ways(const Graph &G, Vertex s, Vertex t,
                                DistanceMap<Length> &distance_s,
                                DistanceMap<Length> &distance_t) {
  using namespace boost;

  auto predecessor = std::vector<Vertex>(num_vertices(G), s);
  auto vertex_id = get(vertex_index, G);

  // Forward step
  dijkstra_shortest_paths(G, s,
                          distance_map(&distance_s[0])
                              .predecessor_map(make_iterator_property_map(
                                  std::begin(predecessor), vertex_id, s)));
  auto edge_list = build_edge_list_from_dijkstra(G, s, t, predecessor);

  // Backward step
  auto rev_G = make_reverse_graph(G);
  dijkstra_shortest_paths(rev_G, t, distance_map(&distance_t[0]));

  if (exists_path_to<Length>(t, distance_s)) {
    return std::make_optional(edge_list);
  } else {
    // In case t could not be found from astar_search and target_found is not
    // thrown, return empty optional
    return std::optional<std::vector<Edge>>{};
  }

} // namespace kspwlo_impl

/**
 * Computes the Dijkstra shortest path from s to t using a
 *        penalty_functor to gather edges weight instead of the Graph's weight
 *        property map
 *
 * @tparam Graph A Boost::PropertyGraph having at least one edge
 *               property with tag boost::edge_weight_t.
 * @tparam PMap The graph's weight property map.
 * @tparam Vertex A vertex_descriptor.
 * @tparam Edge An edge_descriptor.
 * @tparam Length The edge weight type.
 * @param G The graph
 * @param s The source vertex
 * @param t The target vertex
 * @param penalty A penalty_functor
 * @return A vector of the edges of the shortest path from s to t.
 *         An empty optional if t is not reachable from s.
 */
template <typename Graph, typename PMap, typename Vertex = vertex_of_t<Graph>,
          typename Edge = edge_of_t<Graph>,
          typename Length = length_of_t<Graph>>
std::optional<std::vector<Edge>>
dijkstra_shortest_path(const Graph &G, Vertex s, Vertex t,
                       penalty_functor<PMap> &penalty) {
  using namespace boost;

  auto predecessor = std::vector<Vertex>(num_vertices(G), s);
  auto vertex_id = get(vertex_index, G);

  auto weight = make_function_property_map<Edge>(penalty);
  try {
    dijkstra_shortest_paths(
        G, s,
        weight_map(weight)
            .predecessor_map(make_iterator_property_map(std::begin(predecessor),
                                                        vertex_id, s))
            .visitor(make_dijkstra_visitor(
                make_target_visitor(t, on_examine_vertex{}))));
  } catch (target_found tf) {
    auto edge_list = build_edge_list_from_dijkstra(G, s, t, predecessor);
    return std::make_optional(edge_list);
  }

  // In case t could not be found from astar_search and target_found is not
  // thrown, return empty optional
  return std::optional<std::vector<Edge>>{};
}

template <typename Graph, typename PMap, typename AStarHeuristic,
          typename Vertex = vertex_of_t<Graph>,
          typename Edge = edge_of_t<Graph>,
          typename Length = length_of_t<Graph>>
std::optional<std::vector<Edge>>
astar_shortest_path(const Graph &G, Vertex s, Vertex t,
                    penalty_functor<PMap> &penalty,
                    const AStarHeuristic &heuristic) {
  using namespace boost;
  auto predecessor = std::vector<Vertex>(num_vertices(G), s);
  auto vertex_id = get(vertex_index, G);

  auto weight = make_function_property_map<Edge>(penalty);
  try {
    astar_search(G, s, heuristic,
                 predecessor_map(make_iterator_property_map(
                                     std::begin(predecessor), vertex_id, s))
                     .visitor(astar_target_visitor{t})
                     .weight_map(weight));
  } catch (target_found &tf) {
    auto edge_list = build_edge_list_from_dijkstra(G, s, t, predecessor);
    return std::make_optional(edge_list);
  }
  // In case t could not be found from astar_search and target_found is not
  // thrown, return empty optional
  return std::optional<std::vector<Edge>>{};
}

template <typename Graph, typename PMap, typename Vertex = vertex_of_t<Graph>,
          typename Edge = edge_of_t<Graph>,
          typename Length = length_of_t<Graph>>
std::optional<std::vector<Edge>>
bidirectional_dijkstra_shortest_path(const Graph &G, Vertex s, Vertex t,
                                     penalty_functor<PMap> &penalty) {
  using namespace boost;

  auto index = get(vertex_index, G);
  auto predecessor_vec = std::vector<Vertex>(num_vertices(G), s);
  auto predecessor = make_iterator_property_map(predecessor_vec.begin(), index);
  auto distance_vec = std::vector<Length>(num_vertices(G));
  auto distance = make_iterator_property_map(distance_vec.begin(), index);
  auto weight = make_function_property_map<Edge>(penalty);

  auto rev_G = make_reverse_graph(G);
  auto rev_weight_ = reverse_penalty_functor(penalty, G, rev_G);
  using RevEdge = typename boost::graph_traits<
      boost::reverse_graph<Graph>>::edge_descriptor;
  auto rev_weight = make_function_property_map<RevEdge>(rev_weight_);
  auto rev_index = get(vertex_index, rev_G);

  try {
    bidirectional_dijkstra(G, s, t, predecessor, distance, weight, rev_G,
                           rev_weight, rev_index);
  } catch (details::target_not_found &) {
    // In case t could not be found return empty optional
    return std::optional<std::vector<Edge>>{};
  }

  auto edge_list = build_edge_list_from_dijkstra(G, s, t, predecessor);
  return std::make_optional(edge_list);
}

template <typename Graph, typename PMap, typename Vertex = vertex_of_t<Graph>,
          typename Edge = edge_of_t<Graph>>
constexpr std::function<std::optional<std::vector<Edge>>(
    const Graph &, Vertex, Vertex, penalty_functor<PMap> &)>
build_shortest_path_fn(routing_kernels algorithm, const Graph &, const PMap &) {
  switch (algorithm) {
  case routing_kernels::dijkstra:
    return [](const auto &G, auto s, auto t, auto &penalty) {
      return dijkstra_shortest_path(G, s, t, penalty);
    };
  case routing_kernels::bidirectional_dijkstra:
    return [](const auto &G, auto s, auto t, auto &penalty) {
      return bidirectional_dijkstra_shortest_path(G, s, t, penalty);
    };
  default:
    throw std::invalid_argument{
        "Invalid algorithm. Only [dijkstra|bidirectional_dijkstra] allowed."};
  }
}

template <typename Graph, typename PMap, typename AStarHeuristic,
          typename Vertex = vertex_of_t<Graph>,
          typename Edge = edge_of_t<Graph>>
constexpr std::function<std::optional<std::vector<Edge>>(
    const Graph &, Vertex, Vertex, penalty_functor<PMap> &)>
build_shortest_path_fn(routing_kernels algorithm, const Graph &, const PMap &,
                       const AStarHeuristic &heuristic) {
  switch (algorithm) {
  case routing_kernels::astar:
    return [&heuristic](const auto &G, auto s, auto t, auto &penalty) {
      return astar_shortest_path(G, s, t, penalty, heuristic);
    };
  default:
    throw std::invalid_argument{"Invalid algorithm. Only [astar] allowed."};
  }
}

/**
 * Apply penalization step to the candidate path.
 *
 * @pre For each vertex @c v in @p candidate @p distance_s contains the shortest
 *      path distance of @c v from @p s
 * @pre For each vertex @c v in @p candidate @p distance_t contains the shortest
 *      path distance of @c v from @p t
 * @post For each edge @c e in @p candidate, if <tt>penalty_bounds[e] <
 *       bound_limit</tt>, @c e is penalized in @p penalty according to the
 *       following formula: <tt>w(e)_new = w(e) + @p p * w(e)</tt>
 * @post For each edge @c e incoming to a vertex @c u in @p candidate, if
 *       <tt>penalty_bounds[e] < bound_limit</tt>, @c e is penalized in @p
 *       penalty according to the following formula:
 *       <tt>w(e)_new = w(e) + w(e) * (0.1 + @p r * @p distance_t[u] / @p
 *       distance_t[s])</tt>
 * @post For each edge @c e outgoing from a vertex @c v in @p candidate, if
 *       <tt>penalty_bounds[e] < bound_limit</tt>, @c e is penalized in @p
 *       penalty according to the following formula:
 *       <tt>w(e)_new = w(e) + w(e) * (0.1 + @p r * @p distance_s[v] / @p
 *       distance_s[t])</tt>
 *
 * @tparam Graph A Boost::PropertyGraph having at least one edge
 *               property with tag boost::edge_weight_t.
 * @tparam DistanceMap A DistanceMap.
 * @tparam PMap The graph's weight property map.
 * @tparam Vertex A vertex_descriptor.
 * @tparam Edge An edge_descriptor.
 * @tparam Length The edge weight type.
 * @param candidate The candidate path.
 * @param G The graph.
 * @param s The source vertex.
 * @param t The target vertex.
 * @param p The penalty factor for edges in the candidate path.
 * @param r The penalty factor for edges incoming and outgoing to/from vertices
 *          of the candidate path.
 * @param penalty A penalty_functor.
 * @param distance_s A DistanceMap from s.
 * @param distance_t A DistanceMap from t.
 * @param penalty_bounds A PenBoundsMap.
 * @param bound_limit The maximum number of times an edge can be penalized.
 */
template <typename Graph, typename DistanceMap, typename PMap,
          typename Vertex = vertex_of_t<Graph>,
          typename Edge = edge_of_t<Graph>,
          typename Length = length_of_t<Graph>>
void penalize_candidate_path(const std::vector<Edge> &candidate, const Graph &G,
                             Vertex s, Vertex t, double p, double r,
                             penalty_functor<PMap> &penalty,
                             const DistanceMap &distance_s,
                             const DistanceMap &distance_t,
                             PenBoundsMap<Edge> &penalty_bounds,
                             int bound_limit) {
  using namespace boost;

  // Keep track of candidate vertices to exclude them from incoming/outgoing
  // edges update
  auto candidate_vertices = std::unordered_set<Vertex, boost::hash<Vertex>>{};
  auto candidate_edges = std::unordered_set<Edge, boost::hash<Edge>>{};

  for (const auto &e : candidate) {
    candidate_edges.insert(e);

    auto u = source(e, G);
    auto v = target(e, G);
    candidate_vertices.insert(u);
    candidate_vertices.insert(v);
  }

  for (auto &e : candidate_edges) {
    auto u = source(e, G);
    auto v = target(e, G);

    // Check if 'e' is already part of the alternative graph
    if (auto search = penalty_bounds.find(e);
        search != std::end(penalty_bounds)) {
      // If so, penalize only if limit isnt reached
      auto n_updates = search->second;
      if (n_updates < bound_limit) {
        penalty[e] += p * penalty[e];
        ++penalty_bounds[e];
      }
    } else { // Penalize and create a nb_updates counter for e
      penalty[e] += p * penalty[e];
      penalty_bounds.insert({e, 1});
    }

    // Update incoming edges
    for (auto [it, end] = in_edges(u, G); it != end; ++it) {
      auto a = source(*it, G);

      // Incoming edge (a, u) is updated only if 'a' is not part of candidate
      // path
      if (candidate_vertices.find(a) == std::end(candidate_vertices)) {
        // Check if '*it' is already part of the alternative graph
        if (auto search = penalty_bounds.find(*it);
            search != std::end(penalty_bounds)) {
          auto n_updates = search->second;
          // penalize only if limit isnt reached
          if (n_updates < bound_limit) {
            auto closeness = distance_t[u] / distance_t[s];
            auto pen_factor = 0.1 + r * closeness;
            penalty[*it] += pen_factor * penalty[*it];
            ++penalty_bounds[*it];
          }
        } else {
          // Else, just update it
          auto closeness = distance_t[u] / distance_t[s];
          auto pen_factor = 0.1 + r * closeness;
          penalty[*it] += pen_factor * penalty[*it];
        }
      }
    }

    // Update outgoing edges
    for (auto [it, end] = out_edges(v, G); it != end; ++it) {
      auto b = target(*it, G);

      // Outgoing edge (v, b) is updated only if 'b' is not part of candidate
      // path
      if (candidate_vertices.find(b) == std::end(candidate_vertices)) {
        // Check if '*it' is already part of the alternative graph
        if (auto search = penalty_bounds.find(*it);
            search != std::end(penalty_bounds)) {
          auto n_updates = search->second;
          // penalize only if limit isnt reached
          if (n_updates < bound_limit) {
            auto closeness = distance_s[v] / distance_s[t];
            auto pen_factor = 0.1 + r * closeness;
            penalty[*it] += pen_factor * penalty[*it];
            ++penalty_bounds[*it];
          }
        } else {
          // Else, just update it
          auto closeness = distance_s[v] / distance_s[t];
          auto pen_factor = 0.1 + r * closeness;
          penalty[*it] += pen_factor * penalty[*it];
        }
      }
    }
  }
}

template <typename Graph, typename WeightMap, typename MultiPredecessorMap,
          typename RoutingKernel, typename Terminator,
          typename Vertex = vertex_of_t<Graph>>
void penalty(const Graph &G, WeightMap const &original_weight,
             MultiPredecessorMap &predecessors, Vertex s, Vertex t, int k,
             double theta, double p, double r, int max_nb_updates,
             int max_nb_steps, RoutingKernel &routing_kernel,
             Terminator &&terminator) {
  using namespace boost;
  using Edge = typename graph_traits<Graph>::edge_descriptor;
  using Length = typename boost::property_traits<typename boost::property_map<
      Graph, boost::edge_weight_t>::type>::value_type;

  BOOST_CONCEPT_ASSERT((VertexAndEdgeListGraphConcept<Graph>));
  BOOST_CONCEPT_ASSERT((LvaluePropertyMapConcept<WeightMap, Edge>));

  // P_LO set of k paths
  auto resPathsEdges = std::vector<std::vector<Edge>>{};
  auto resEdges = std::vector<std::unordered_set<Edge, boost::hash<Edge>>>{};

  // P_LO set of k paths
  auto resPaths = std::vector<Path<Graph>>{};
  // Make a local weight map to avoid modifying existing graph.
  auto pen_fctor = details::penalty_functor{original_weight};

  // Compute shortest path from s to t
  auto distance_s = std::vector<Length>(num_vertices(G));
  auto distance_t = std::vector<Length>(num_vertices(G));
  auto sp = dijkstra_shortest_path_two_ways(G, s, t, distance_s, distance_t);
  if (!sp) {
    auto oss = std::ostringstream{};
    oss << "Vertex " << t << " is unreachable from " << s;
    throw details::target_not_found{oss.str()};
  }

  // P_LO <-- {shortest path p_0(s, t)};
  resPathsEdges.push_back(*sp);
  resEdges.emplace_back(sp->begin(), sp->end());

  // If we need the shortest path only
  if (k == 1) {
    fill_multi_predecessor(resPathsEdges.begin(), resPathsEdges.end(), G,
                           predecessors);
    return;
  }

  // Initialize map for penalty bounds
  auto penalty_bounds = std::unordered_map<Edge, int, boost::hash<Edge>>{};

  // Penalize sp edges
  penalize_candidate_path(*sp, G, s, t, p, r, pen_fctor, distance_s, distance_t,
                          penalty_bounds, max_nb_updates);

  int step = 0;
  using Index = std::size_t;
  while (resPathsEdges.size() < static_cast<Index>(k) && step < max_nb_steps) {
    // The remainder code is the hot part of the algorithm. So we check here
    // if the algorithm should terminate
    if (terminator.should_stop()) {
      throw terminator_stop_error{
          "Penalty terminated before completing due to a Terminator. Please "
          "discard partial output."};
    }

    auto p_tmp = routing_kernel(G, s, t, pen_fctor);

    // Penalize p_tmp edges
    penalize_candidate_path(*p_tmp, G, s, t, p, r, pen_fctor, distance_s,
                            distance_t, penalty_bounds, max_nb_updates);
    ++step;

    // If p_tmp is sufficiently dissimilar to other alternative paths, accept it
    bool is_valid_path = true;
    for (const auto &alt_path : resEdges) {
      if (compute_similarity(*p_tmp, alt_path, original_weight) > theta) {
        is_valid_path = false;
        break;
      }
    }

    if (is_valid_path) {
      resPathsEdges.push_back(*p_tmp);
      resEdges.emplace_back(p_tmp->begin(), p_tmp->end());
    }
  }

  // Beforer returning, populate predecessors map
  fill_multi_predecessor(resPathsEdges.begin(), resPathsEdges.end(), G,
                         predecessors);
}
} // namespace details
} // namespace arlib

#endif