// L-40 MCS 360 Fri 1 Dec 2017 : minimum_spanning_tree.cpp

// Runs Prim's algorithm on a randomly generated graph
// to compute the minimum spanning tree.

#include <iostream>
#include <iomanip>
#include <cstdlib>
#include <vector>
#include <limits>
#include <utility>
#include <map>
#include <set>

using namespace std;

void allocate_double_matrix ( vector< vector<double> >& mat, size_t dim );
/*
   Allocates space in mat for an adjacency matrix of dimension dim. */

int loaded_coin ( double p );
/*
   A fair coin will turn out as many 0 as 1 if p = 0.5,
   as p is the probability of 0. */

void random_undirected_weighted_graph
 ( vector< vector<double> >& mat, double prb );
/*
   Fills up the adjacency matrix mat for a random undirected graph.
   The value of prb is the probability that there is an edge.
   The weights are random doubles between 0.1 and 9.9. */

void write_double_matrix ( vector< vector<double> >& mat );
/*
   Writes the matrix mat to screen. */

void write_data ( vector<double>& dis, vector<size_t>& prv );
/*
   Writes the distances dis and predecessors prv. */

void write_sets ( set<size_t>& done, set<size_t>& todo );
/*
   Writes the sets done and todo. */

size_t idx_min_dis ( set<size_t> s, vector<double> d );
/*
   Returns the index of the element in s which has the smallest
   value for its distance in the corresponding vector d.
   Returns d.size() if there is no index with distance < oo. */

void initialize
 ( vector< vector<double> >& mat, size_t start,
   vector<double>& dis, vector<size_t>& prv,
   set<size_t>& done, set<size_t>& todo, bool verbose );
/*
   Given the adjacency matrix mat for a directed weighted graph
   and a vertex start, initializes the vectors dis and prv,
   and initializes the sets done and todo.
   If verbose, then dis and prv are printed. */

void minimum_spanning_tree
 ( vector< vector<double> >& mat, size_t start,
   vector<double>& dis, vector<size_t>& prv,
   multimap<size_t, size_t>& mintree,  bool verbose );
/*
   Given the adjacency matrix mat for a undirected weighted graph
   and a vertex start, returns two vectors dis, prv, and the set
   of edges in the minimum spanning tree.
   On return, dis[k] the distance of the shortest edge to vertex k,
   prv[k] is the source of the shortest edge to vertex k, and
   mintree contains edges in the minimum spanning tree,
   the index in the multiple key map is the source vertex of each edge.
   If verbose, then intermediate values of dis and prv are printed. */

void write_tree
 ( multimap<size_t, size_t>& tree, size_t root, size_t level );
/*
   Given in tree is the representation of a tree, stored in 
   a multiple key map.  For each edge there is an entry in the map.
   The source of each edge is the key, the value is the destination
   of each edge.  The current root of the tree is the parameter root.
   The depth of the current root is given by level. */

double cost_tree
 ( vector< vector<double> >& mat,
   multimap<size_t, size_t>& tree, size_t root );
/*
   Returns the cost of the spanning tree for the weights defined
   in the adjacency matrix mat.
   Given in tree is the representation of a tree, stored in 
   a multiple key map.  For each edge there is an entry in the map.
   The source of each edge is the key, the value is the destination
   of each edge.  The current root of the tree is the parameter root.
   The depth of the current root is given by level. */

void write_results
 ( vector< vector<double> >& mat, size_t start,
   vector<double>& dis, vector<size_t>& prv, multimap<size_t, size_t>& tree );
/*
   Given the adjacency matrix mat for a directed weighted graph
   and a vertex start, writes the results of the shortest_paths:
   dis[k] the shortest distance of vertex start to vertex k, and
   prv[k] is the previous vertex in the shortest path from start to k. */

int main ( void )
{
   size_t seed = time(NULL);
   cout << "The seed : " << seed << endl;
   srand(seed); // srand(time(NULL));

   cout << "Give the number of vertices : ";
   size_t n; cin >> n;

   cout << "Give a probability in [0.0, 1.0] : ";
   double p; cin >> p;

   vector< vector<double> > mat;
   allocate_double_matrix(mat, n);
   random_undirected_weighted_graph(mat, p);
   write_double_matrix(mat);

   cout << "Give the vertex start : ";
   size_t s; cin >> s; 

   vector<double> distances;
   for(size_t i=0; i<mat.size(); i++) distances.push_back(0.0);
   vector<size_t> predecessors;
   for(size_t i=0; i<mat.size(); i++) predecessors.push_back(0);

   multimap<size_t, size_t> tree;

   minimum_spanning_tree(mat, s, distances, predecessors, tree, true);

   write_results(mat, s, distances, predecessors, tree);

   return 0;
}

void allocate_double_matrix ( vector< vector<double> >& mat, size_t dim )
{
   for(size_t i=0; i<dim; i++)
   {
      vector<double> row;

      for(size_t j=0; j<dim; j++) row.push_back(0.0);

      mat.push_back(row);
   }
}

int loaded_coin ( double p )
{
   double r = ((double) rand())/RAND_MAX;

   return (r <= p ? 0 : 1);
}

void random_undirected_weighted_graph
 ( vector< vector<double> >& mat, double prb )
{
   const size_t dim = mat.size();

   for(size_t i=0; i<dim; i++)
   {
      for(size_t j=0; j<i; j++)  // copy upper half
         mat[i][j] = mat[j][i];
      for(size_t j=i; j<dim; j++)  // place edges at random
      {
         size_t rnd = loaded_coin(prb);  // edge or not
         if(rnd == 0)
            mat[i][j] = 0.0;       // no edge => zero weight
         else
         {
            rnd = 1 + rand() % 99; // random weight 
            mat[i][j] = ((double) rnd)/10.0; // in [0.1, 9.9]
         }
      }
   }
}

void write_double_matrix ( vector< vector<double> >& mat )
{
   const size_t dim = mat.size();

   cout << fixed << setprecision(1);
   
   for(size_t i=0; i<dim; i++)
   {
      for(size_t j=0; j<dim; j++)
         cout << " " << mat[i][j];
      cout << endl;
   }  
}

void write_data ( vector<double>& dis, vector<size_t>& prv )
{
   cout << "The distances :" << endl;
   for(size_t i=0; i<dis.size(); i++)
      cout << " " << dis[i];
   cout << endl;

   cout << "The predecessors :" << endl;
   for(size_t i=0; i<prv.size(); i++)
      cout << " " << prv[i];
   cout << endl;
}

void write_sets ( set<size_t>& done, set<size_t>& todo )
{
   cout << "The set done :";
   for(set<size_t>::const_iterator i=done.begin(); i != done.end(); i++)
      cout << " " << *i;
   cout << endl;

   cout << "The set todo :";
   for(set<size_t>::const_iterator i=todo.begin(); i != todo.end(); i++)
      cout << " " << *i;
   cout << endl;
}

size_t idx_min_dis ( set<size_t> s, vector<double> d )
{
   size_t result = d.size();
   const double oo = numeric_limits<double>::infinity();
   double mindis = oo;

   for(set<size_t>::const_iterator i=s.begin(); i != s.end(); i++)
   {
      if(d[*i] < mindis)
      {
         mindis = d[*i];
         result = *i;
      }
   }

   return result;
}

void initialize
 ( vector< vector<double> >& mat, size_t start,
   vector<double>& dis, vector<size_t>& prv,
   set<size_t>& done, set<size_t>& todo, bool verbose )
{
   const size_t dim = mat.size();
   const double oo = numeric_limits<double>::infinity();

   for(size_t i=0; i<dim; i++)
   {
      dis[i] = (mat[start][i] == 0.0 ? oo : mat[start][i]);
      prv[i] = start;
   }
   done.insert(start);
   for(size_t i=0; i<dim; i++)
      if(i != start) todo.insert(i);

   if(verbose)
   {
      write_data(dis, prv);
      write_sets(done, todo);
   }
}

void minimum_spanning_tree
 ( vector< vector<double> >& mat, size_t start,
   vector<double>& dis, vector<size_t>& prv,
   multimap<size_t, size_t>& mintree,  bool verbose )
{
   set<size_t> done; // vertices already processed
   set<size_t> todo; // vertices left to process
   initialize(mat, start, dis, prv, done, todo, verbose);

   while(not todo.empty())
   {
      size_t vtx = idx_min_dis(todo, dis); // closest vertex
      if(vtx == dis.size())
      {
         if(verbose) cout << "No distance < oo." << endl;
         break;
      }
      done.insert(vtx); todo.erase(vtx);   // remove from todo
      if(verbose) write_sets(done,todo);

      pair<size_t, size_t> edge;
      edge.first = prv[vtx];
      edge.second = vtx;
      mintree.insert(edge);

      for(set<size_t>::const_iterator i=todo.begin(); i!=todo.end(); i++)
         if(mat[vtx][*i] != 0.0)
            if(dis[vtx] + mat[vtx][*i] < dis[*i])
            {
               dis[*i] = mat[vtx][*i];
               prv[*i] = vtx;
            }
      if(verbose) write_data(dis, prv);
   }
}

void write_tree
 ( multimap<size_t, size_t>& tree, size_t root, size_t level )
{
   for(size_t k=0; k<level; k++) cout << "  ";
   cout << root << endl;
   
   for(multimap<size_t, size_t>::const_iterator it = tree.begin();
       it != tree.end(); it++)
      if(it->first == root) write_tree(tree, it->second, level+1);
}

double cost_tree
 ( vector< vector<double> >& mat,
   multimap<size_t, size_t>& tree, size_t root )
{
   double result = 0.0;

   for(multimap<size_t, size_t>::const_iterator it = tree.begin();
       it != tree.end(); it++)
      if(it->first == root)
      {
         double cost_edge = mat[root][it->second];
         double cost_subtree = cost_tree(mat, tree, it->second);

         result = result + cost_edge + cost_subtree;
      }
   return result;
}

void write_results
 ( vector< vector<double> >& mat, size_t start,
   vector<double>& dis, vector<size_t>& prv, multimap<size_t, size_t>& tree )
{
   cout << "The adjacency matrix :" << endl;
   write_double_matrix(mat);
   cout << "The edges in the minimum spanning tree :" << endl;
   for(multimap<size_t, size_t>::const_iterator i=tree.begin();
       i!=tree.end(); i++)
      cout << "(" << i->first << "," << i->second << ")";
   cout << endl;
   cout << "The minimum spanning tree :" << endl;
   write_tree(tree, start, 0);
   cout << "The cost of the spanning tree : ";
   cout << cost_tree(mat, tree, start);
   cout << endl;
   double sumdis = 0.0;
   for(size_t i=0; i<dis.size(); i++)
      if(i != start) sumdis = sumdis + dis[i];
   cout << "Sum of distances of shortest edges : " << sumdis << endl;
}