class Vertex {
public:
    Vertex() {
        bVisited = false;
    }
    int data = 0;
    EdgeList edgeList;
    bool bVisited = false;
};

struct Edge {
    int cost;
    boost::weak_ptr<Vertex> srcVertex;
    boost::weak_ptr<Vertex> dstVertex;
};

typedef std::list<boost::shared_ptr<Vertex>> VerticesList;
typedef std::map<int, VerticesList::const_iterator> DataVertexMap;
class Graph {

    VerticesList _verticesList;
    DataVertexMap _dataVertexMap;

    //TODO: Need some design to have this function internal only and not exposed to client
    boost::shared_ptr<Vertex> addAndGetVertex(int data);
public:
    Graph();
    ~Graph();

    bool isEmpty() const;
    //We don't check for duplicate vertex
    void addVertex(int data);
    void addEdge(int srcData, int dstData, int cost);
    int getCostForEdge(int srcData, int dstData);
    void displayGraph();
    void dfsTraversal();
    void bfsTraversal();

    void findShortestPath(int srcData, int dstData);
    void kruskalMST();
    void primMST();
};

void Graph::findShortestPath(int srcData, int dstData) {
    DataVertexMap::const_iterator srcIter = _dataVertexMap.find(srcData);
    DataVertexMap::const_iterator dstIter = _dataVertexMap.find(dstData);
    if(srcIter == _dataVertexMap.end() || dstIter == _dataVertexMap.end())
        assert(false);

    boost::shared_ptr<Vertex> srcVertex(*(srcIter->second));
    boost::shared_ptr<Vertex> dstVertex(*(dstIter->second));

    struct VertexInfo {
        bool bGotMinDistance = false;
        boost::weak_ptr<Vertex> precedingVertex;
        int distanceFromSource = INFINITY;
    };

    boost::shared_ptr<Vertex> currentVertex = srcVertex;

    typedef std::map<boost::shared_ptr<Vertex>, VertexInfo> VertexInfoMap;
    VertexInfoMap vertexInfoMap;
    for(boost::shared_ptr<Vertex> vertex : _verticesList) {
        VertexInfo info;
        if(vertex.get() == srcVertex.get()) {
            info.bGotMinDistance = true;
            info.distanceFromSource = 0;
            info.precedingVertex = boost::weak_ptr<Vertex> ();
        }
        vertexInfoMap.insert(std::make_pair(vertex, info));
    }


    while(currentVertex.get() != dstVertex.get()) {
        VertexInfo& curVertexInfo = vertexInfoMap[currentVertex];
        curVertexInfo.bGotMinDistance = true;
        boost::shared_ptr<Vertex> minDistVertex = currentVertex;
        int curMinDist  = INFINITY;
        for(auto iter = vertexInfoMap.begin(); iter != vertexInfoMap.end(); ++iter) {
            boost::shared_ptr<Vertex> vertex = iter->first;
            VertexInfo& vertexInfo = iter->second;
            if(!vertexInfo.bGotMinDistance) {
                int distFromCurToVertex = getCostForEdge(currentVertex->data, vertex->data);
                if(vertexInfo.distanceFromSource > distFromCurToVertex + curVertexInfo.distanceFromSource) {
                    vertexInfo.distanceFromSource = distFromCurToVertex + curVertexInfo.distanceFromSource;
                    vertexInfo.precedingVertex = currentVertex;
                }

                if(vertexInfo.distanceFromSource < curMinDist) {
                    curMinDist = vertexInfo.distanceFromSource;
                    minDistVertex = vertex;
                }
            }
        }
        currentVertex = minDistVertex;
    }

class CompareEdge {
public:
    bool operator() (const boost::shared_ptr<Edge> & e1, const boost::shared_ptr<Edge> & e2) const{
        if(e1->cost == e2->cost)
            return e1.get() < e2.get();  //similar to default comparator less<T>

        return e1->cost < e2->cost;
    }
};

typedef std::map<boost::shared_ptr<Vertex>, boost::shared_ptr<Vertex>> ParentMap;

// A utility function to find the subset of an element i
boost::shared_ptr<Vertex> find(const ParentMap& parentMap, const boost::shared_ptr<Vertex>& vertex)
{
    ParentMap::const_iterator iter = parentMap.find(vertex);
    if (!(iter->second).get())
        return vertex;
    return find(parentMap, iter->second);
}

// A utility function to do union of two subsets
void Union(ParentMap& parentMap, const boost::shared_ptr<Vertex>& x, const boost::shared_ptr<Vertex>& y)
{
    const boost::shared_ptr<Vertex> xset = find(parentMap, x);
    const boost::shared_ptr<Vertex> yset = find(parentMap, y);
    parentMap[x] = y;
}


// The main function to check whether a given graph contains cycle or not
bool isCycle(const VerticesList& verticesList, const EdgeList& allEdgeList)
{
    // Allocate memory for creating V subsets
    ParentMap parent;

    for(boost::shared_ptr<Vertex> vertex : verticesList) {
        parent.insert(std::make_pair(vertex, nullptr));
    }

    // Iterate through all edges of graph, find subset of both
    // vertices of every edge, if both subsets are same, then there is
    // cycle in graph.
    for(boost::shared_ptr<Edge> edge : allEdgeList)
    {
        boost::shared_ptr<Vertex> x = find(parent, (edge->srcVertex).lock());
        boost::shared_ptr<Vertex> y = find(parent, (edge->dstVertex).lock());

        if (x.get() == y.get())
            return true;

        Union(parent, x, y);
    }
    return false;
}

void Graph::kruskalMST() {
    //set to hold all edges in increasing order of edge cost
    std::set<boost::shared_ptr<Edge>, CompareEdge> allEdgeList;
    for(boost::shared_ptr<Vertex> vertex : _verticesList) {
        EdgeList& vertexEdgeList = vertex->edgeList;
        allEdgeList.insert(vertexEdgeList.begin(), vertexEdgeList.end());
    }

    EdgeList mstEdgeList;
    for(boost::shared_ptr<Edge> edge : allEdgeList) {
        if(mstEdgeList.size() == _verticesList.size() - 1) break;
        mstEdgeList.push_front(edge);
        if(isCycle(_verticesList, mstEdgeList)) {
            //Remove the edge from MST
            mstEdgeList.pop_front();
        }
    }

    // print the contents of result[] to display the built MST
    printf("Following are the edges in the constructed MSTn");
    for (boost::shared_ptr<Edge> edge : mstEdgeList)
        printf("%d -- %d == %dn", (edge->srcVertex).lock()->data, (edge->dstVertex).lock()->data,
               edge->cost);
    return;
}

//return min cost edge
 boost::shared_ptr<Edge> getMinEdge(const EdgeList& edgeList) {
    int min = INFINITY;
    boost::shared_ptr<Edge> minEdge;
    for(boost::shared_ptr<Edge> edge : edgeList) {
        if(edge->cost < min) {
            min = edge->cost;
            minEdge = edge;
        }
    }

    return minEdge;
}

void Graph::primMST() {
    EdgeList mstEdgeList;
    EdgeList potentialEdgeList;

    //insert edges for first vertex
    boost::shared_ptr<Vertex> firstVertex = _verticesList.front();
    potentialEdgeList.insert(potentialEdgeList.begin(), firstVertex->edgeList.begin(), firstVertex->edgeList.end());

    while(mstEdgeList.size() < _verticesList.size() - 1) {
        boost::shared_ptr<Edge> minEdge = getMinEdge(potentialEdgeList);
        mstEdgeList.push_front(minEdge);
        potentialEdgeList.remove(minEdge);
        //Mark the vertices visited so that edge having them as destination are not again included
        (minEdge->srcVertex).lock()->bVisited = true;
        (minEdge->dstVertex).lock()->bVisited = true;
        EdgeList& dstVertexEdgeList = (minEdge->dstVertex).lock()->edgeList;
        for(boost::shared_ptr<Edge> edge : dstVertexEdgeList) {
            if(!(edge->dstVertex).lock()->bVisited)
                potentialEdgeList.push_front(edge);
        }
    }

    //TODO: set all bVisited false

    // print the contents of result[] to display the built MST
    printf("Following are the edges in the constructed MSTn");
    for (boost::shared_ptr<Edge> edge : mstEdgeList)
        printf("%d -- %d == %dn", (edge->srcVertex).lock()->data, (edge->dstVertex).lock()->data,
               edge->cost);
    return;
}