// -*- c++ -*- //
#ifndef HUGO_KRUSKAL_H
#define HUGO_KRUSKAL_H

#include <algorithm>
#include <unionfind.h>
#include <for_each_macros.h>

namespace hugo {

  /// Kruskal's algorithm to compute a minimum cost tree

  template <typename Graph, typename InputEdgeOrder, typename OutBoolMap>
  typename InputEdgeOrder::ValueType
  Kruskal(Graph const& G, InputEdgeOrder const& edges, 
	  OutBoolMap& out_map)
  {
    typedef typename InputEdgeOrder::ValueType EdgeCost;
    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 InputEdgeOrder::const_iterator p = edges.begin(); 
	 p!=edges.end(); ++p ) {
      if ( uf.join(G.head(edges.first(p)),
			     G.tail(edges.first(p))) ) {
	out_map.set(edges.first(p), true);
	tot_cost += edges.second(p);
      }
      else {
	out_map.set(edges.first(p), false);
      }
    }
    return tot_cost;
  }

  /* A work-around for running Kruskal with const-reference bool maps... */

  template<typename Map>
  class NonConstMapWr {
    const Map &m;
  public:
    typedef typename Map::ValueType ValueType;

    NonConstMapWr(const Map &_m) : m(_m) {}

    template<typename KeyType>
    void set(KeyType const& k, ValueType const &v) const { m.set(k,v); }
  };

  template <typename Graph, typename InputEdgeOrder, typename OutBoolMap>
  inline
  typename InputEdgeOrder::ValueType
  Kruskal(Graph const& G, InputEdgeOrder const& edges, 
	  OutBoolMap const& out_map)
  {
    NonConstMapWr<OutBoolMap> map_wr(out_map);
    return Kruskal(G, edges, map_wr);
  }  

  
  /* ** ** Output-objektumok: egyszeruen extra bool mapek ** ** */
  
  /// A writable bool-map that makes a sequence of "true" keys

  /// A writable bool-map that creates a sequence out of keys that receive
  /// the value "true".
  /// \warning Not a regular property map, as it doesn't know its KeyType

  template<typename Iterator>
  class SequenceOutput {
    mutable Iterator it;

  public:
    typedef bool ValueType;

    SequenceOutput(Iterator const &_it) : it(_it) {}

    template<typename KeyType>
    void set(KeyType const& k, bool v) const { if(v) {*it=k; ++it;} }
  };

  template<typename Iterator>
  inline
  SequenceOutput<Iterator>
  makeSequenceOutput(Iterator it) {
    return SequenceOutput<Iterator>(it);
  }

  /* ** ** InputSource -ok ** ** */

  template<typename Key, typename Val>
  class KruskalPairVec : public std::vector< std::pair<Key,Val> > {

  public:
    typedef std::vector< std::pair<Key,Val> > Parent;
    typedef Key KeyType;
    typedef Val ValueType;
    typedef std::pair<Key,Val> PairType;

    typedef typename Parent::iterator iterator;
    typedef typename Parent::const_iterator const_iterator;

  private:
    class comparePair {
    public:
      bool operator()(PairType const& a, PairType const& b) {
	return a.second < b.second;
      }
    };

  public:

    // FIXME: kell ez?
    // KruskalPairVec(Parent const& p) : Parent(p) {}
    
    void sort() {
      std::sort(begin(), end(), comparePair());
    }

    // FIXME: nem nagyon illik ez ide...
    template<typename Graph, typename Map>
    KruskalPairVec(Graph const& G, Map const& m) {
      typedef typename Graph::EdgeIt EdgeIt;

      clear();
      FOR_EACH_LOC(EdgeIt, e, G) {
      // for (EdgeIt e=G.template first<EdgeIt>(); G.valid(e); G.next(e)) {
	push_back(make_pair(e, m[e]));
      }
      sort();
    }

    Key const& first(const_iterator i) const { return i->first; }
    Key& first(iterator i) { return i->first; }

    Val const& second(const_iterator i) const { return i->second; }
    Val& second(iterator i) { return i->second; }
  };


  template <typename Map>
  class KruskalMapVec : public std::vector<typename Map::KeyType> {
  public:

    typedef typename Map::KeyType KeyType;
    typedef typename Map::ValueType ValueType;

    typedef typename std::vector<KeyType> Parent;
    typedef typename Parent::iterator iterator;
    typedef typename Parent::const_iterator const_iterator;

  private:

    const Map &m;

    class compareKeys {
      const Map &m;
    public:
      compareKeys(Map const &_m) : m(_m) {}
      bool operator()(KeyType const& a, KeyType const& b) {
	return m[a] < m[b];
      }
    };

  public:

    KruskalMapVec(Map const& _m) : m(_m) {}

    void sort() {
      std::sort(begin(), end(), compareKeys(m));
    }

    // FIXME: nem nagyon illik ez ide...
    template<typename Graph>
    KruskalMapVec(Graph const& G, Map const& _m) : m(_m) {
      typedef typename Graph::EdgeIt EdgeIt;

      clear();
      FOR_EACH_LOC(EdgeIt, e, G) {
      // for (EdgeIt e=G.template first<EdgeIt>(); G.valid(e); G.next(e)) {
	push_back(e);
      }
      sort();
    }

    KeyType const& first(const_iterator i) const { return *i; }
    KeyType& first(iterator i) { return *i; }

    ValueType const& second(const_iterator i) const { return m[*i]; }
    ValueType& second(iterator i) { return m[*i]; }
  };

  /* ** ** Wrapper fuggvenyek ** ** */


  /// \brief Wrapper to Kruskal().
  /// Input is from an EdgeMap, output is a plain boolmap.
  template <typename Graph, typename EdgeCostMap, typename RetEdgeBoolMap>
  inline
  typename EdgeCostMap::ValueType
  Kruskal_EdgeCostMapIn_BoolMapOut(Graph const& G,
				   EdgeCostMap const& edge_costs,
				   RetEdgeBoolMap &ret_bool_map) {

    typedef KruskalPairVec<typename Graph::Edge, typename EdgeCostMap::ValueType>
      InputVec;
    InputVec iv(G, edge_costs);

    return Kruskal(G, iv, ret_bool_map);
  }


  /// \brief Wrapper to Kruskal().
  /// Input is from an EdgeMap, output is to a sequence.
  template <typename Graph, typename EdgeCostMap, typename RetIterator>
  inline
  typename EdgeCostMap::ValueType
  Kruskal_EdgeCostMapIn_IteratorOut(Graph const& G,
				    EdgeCostMap const& edge_costs,
				    RetIterator ret_iterator) {
    typedef SequenceOutput<RetIterator> OutMap;
    OutMap out(ret_iterator);

    typedef KruskalPairVec<typename Graph::Edge, typename EdgeCostMap::ValueType>
      InputVec;
    InputVec iv(G, edge_costs);

    return Kruskal(G, iv, out);
  }


} //namespace hugo

#endif //HUGO_KRUSKAL_H