Untitled

#include "color_walk.hpp"

using namespace cs251;

std::vector<std::pair<char, int>> color_walk::calculate(const graph& g, const handle startHandle) {
	// TODO: Implement color walk	
	//throw std::logic_error("color_walk::" + std::string(__FUNCTION__) + "() not implemented");
    std::unordered_map<handle, std::vector<handle>> colorVertices;
        graph transformedGraph;

        for (const auto& vertex : g.m_vertices) {
            for (const auto& edge : vertex.m_edges) {
                handle originalVertex = vertex.m_handle;
                handle colorVertex = transformedGraph.add_vertex();
                transformedGraph.add_edge(originalVertex, colorVertex, edge.m_weight);
                colorVertices[originalVertex].push_back(colorVertex);
            }
        }

        // Step 2: Initialize distances and colors for each vertex in G'
        std::unordered_map<handle, std::vector<int>> distances;
        std::unordered_map<handle, std::vector<char>> colors;

        for (const auto& vertex : transformedGraph.m_vertices) {
            distances[vertex.m_handle].resize(3, std::numeric_limits<int>::max());
            colors[vertex.m_handle].resize(3, '-');

            if (vertex.m_handle == startHandle) {
                distances[vertex.m_handle][0] = 0; // Starting color '-'
            }
        }

        // Step 3: Perform Dijkstra's algorithm on G' starting from startHandle's color vertex
        std::priority_queue<std::pair<int, handle>, std::vector<std::pair<int, handle>>, std::greater<>> pq;
        pq.push({ 0, startHandle });

        while (!pq.empty()) {
            auto [dist, u] = pq.top();
            pq.pop();

            for (int i = 0; i < 3; ++i) {
                for (const auto& edge : transformedGraph.m_vertices[u].m_edges) {
                    int v = edge.m_destinationHandle;
                    int weight = edge.m_weight;

                    if (distances[v][i] > dist + weight) {
                        distances[v][i] = dist + weight;
                        colors[v][i] = colors[u][i];

                        pq.push({ distances[v][i], v });
                    }
                }
            }
        }

        // Step 4: Construct and return the result based on the shortest paths found
        std::vector<std::pair<char, int>> result;

        for (const auto& vertex : transformedGraph.m_vertices) {
            for (int i = 0; i < 3; ++i) {
                if (colors[vertex.m_handle][i] != '-') {
                    result.push_back({ colors[vertex.m_handle][i], distances[vertex.m_handle][i] });
                    break; // Only consider the shortest path color and distance
                }
            }
        }

        return result;
    }
}