// -*- c++ -*- //

#ifndef HUGO_KRUSKAL_H

#define HUGO_KRUSKAL_H

#include "unionfind.h"

#include <algorithm>

namespace hugo {

  template <typename Graph, typename EdgeCostMap, typename EdgeBoolMap>

  class MinCostTreeKruskal

  {

    typedef typename Graph::EdgeIt EdgeIt;

    typedef typename Graph::Edge Edge;

  public:

    typedef typename EdgeCostMap::ValueType EdgeCost;

    typedef std::pair<Edge, EdgeCost> EdgeCostPair;

    typedef std::vector<EdgeCostPair> EdgeCostVector;

  private:

    Graph const &G;

    EdgeCostMap const &costmap;

    EdgeBoolMap& treemap;

    //template <typename EdgeCostPair>

    class comparePair {

    public:

      bool operator()(EdgeCostPair const& a, EdgeCostPair const& b) {

	return a.second < b.second;

      }

    };

  public:

    MinCostTreeKruskal(Graph const& _G, EdgeCostMap const& _costmap, 

		       EdgeBoolMap& _treemap) : G(_G), costmap(_costmap), 

						treemap(_treemap) {}

    EdgeCost run() 

    {

     EdgeCostVector rank;

     for (EdgeIt e=G.template first<EdgeIt>(); G.valid(e); G.next(e)) {

	rank.push_back(make_pair(e, costmap.get(e)));

      }

      std::sort(rank.begin(), rank.end(), comparePair());

      return run(rank);

    }

    EdgeCost run(EdgeCostVector const &rank)

    {

      typedef typename Graph::NodeMap<int> NodeIntMap;

      typedef typename Graph::Node Node;

      NodeIntMap comp_map(G, -1);

      UnionFind<Node,NodeIntMap> uf(comp_map); 

      EdgeCost tot_cost = 0;

      for (typename EdgeCostVector::const_iterator p = rank.begin(); 

	   p!=rank.end(); ++p ) {

	if ( uf.joinComponents(G.head(p->first), G.tail(p->first)) ) {

	  treemap.set(p->first, true);

	  tot_cost += p->second;

	}

	else {

	  treemap.set(p->first, false);

	}

      }

      return tot_cost;

    }

  };

} //namespace hugo

#endif //HUGO_KRUSKAL_H