#include "costs.h"
#include "trailblazer.h"

#include <stack>
#include <queue>
#include <vector>
#include <algorithm>
#include <limits>
#include "pqueue.h"

using namespace std;

vector<Node*> path_from_end(Node* end, bool color_green = false) {
    if (!end->previous) {
        return {};
    }
    Node* cur = end;
    vector<Node*> path;
    while (cur) {
        if (color_green) {
            cur->setColor(GREEN);
        }
        path.push_back(cur);
        cur = cur->previous;
    }
    reverse(path.begin(), path.end());
    return path;
}

template <typename T>
void push_back_from(vector<T>& to, const vector<T>& from) {
    for (const auto& el : from) {
        to.push_back(el);
    }
}

vector<Node*> dfs_helper(BasicGraph& graph, Vertex* start, Vertex* end) {
    start->setColor(GREEN);
    vector<Node*> path = { start };
    if (start == end) {
        return path;
    }
    start->visited = true;
    for (const auto& arc : start->arcs) {
        if (!arc->finish->visited) {
            vector<Node*> neighbour_path = dfs_helper(graph, arc->finish, end);
            if (neighbour_path.size() == 0) {
                continue;
            }
            vector<Node*> new_path(path);
            push_back_from(new_path, neighbour_path);
            return new_path;
        }
    }
    start->setColor(GRAY);
    return {};
}

vector<Node*> depthFirstSearch(BasicGraph& graph, Vertex* start, Vertex* end) {
    graph.resetData();
    return dfs_helper(graph, start, end);
}

vector<Node *> breadthFirstSearch(BasicGraph& graph, Vertex* start, Vertex* end) {
    graph.resetData();

    queue<Node*> q;
    q.push(start);
    start->visited = true;
    while (q.size() > 0) {
        Node* cur = q.front();
        q.pop();
        cur->setColor(GREEN);
        if (cur == end) {
            break;
        }
        for (const auto& arc : cur->arcs) {
            if (!arc->finish->visited) {
                arc->finish->visited = true;
                arc->finish->previous = cur;
                arc->finish->setColor(YELLOW);
                q.push(arc->finish);
            }
        }
    }
    return path_from_end(end);
}

vector<Node*> dijkstrasAlgorithm(BasicGraph& graph, Vertex* start, Vertex* end) {
    graph.resetData();

    PriorityQueue<Node *> q;
    q.enqueue(start, 0.0);
    start->visited = true;

    while (q.size() > 0) {
        Node* cur = q.dequeue();
        cur->visited = true;
        cur->setColor(GREEN);
        if (cur == end) {
            break;
        }
        for (const auto& arc : cur->arcs) {
            if (!arc->finish->visited) {
                double new_dist = arc->start->cost + arc->cost;
                if (arc->finish->cost == 0.0) {
                    arc->finish->previous = cur;
                    arc->finish->cost = new_dist;
                    arc->finish->setColor(YELLOW);
                    q.enqueue(arc->finish, new_dist);
                }
                else if (new_dist < arc->finish->cost) {
                    arc->finish->previous = cur;
                    arc->finish->cost = new_dist;
                    q.changePriority(arc->finish, new_dist);
                }
            }
        }
    }
    return path_from_end(end);
}

vector<Node*> aStar(BasicGraph& graph, Vertex* start, Vertex* end) {
    graph.resetData();

    PriorityQueue<Node *> q;
    q.enqueue(start, 0.0);
    start->visited = true;

    while (q.size() > 0) {
        Node* cur = q.dequeue();
        cur->visited = true;
        cur->setColor(GREEN);
        if (cur == end) {
            break;
        }
        for (const auto& arc : cur->arcs) {
            if (!arc->finish->visited) {
                double new_dist = arc->start->cost + arc->cost;
                if (arc->finish->cost == 0.0) {
                    arc->finish->previous = cur;
                    arc->finish->cost = new_dist;
                    arc->finish->setColor(YELLOW);
                    q.enqueue(arc->finish, new_dist + arc->finish->heuristic(end));
                }
                else if (new_dist < arc->finish->cost) {
                    arc->finish->previous = cur;
                    arc->finish->cost = new_dist;
                    q.changePriority(arc->finish, new_dist + arc->finish->heuristic(end));
                }
            }
        }
    }
    return path_from_end(end);
}