#include "graph.h"
#include "test_runner.h"

#include <vector>
#include <string>
#include <algorithm>

using namespace std;

#ifdef LOCAL
#define dbg(x) cerr << #x << " : " << x << endl
#else
#define dbg(x)
#endif

GraphManager::GraphManager(double wait_, double speed_) : wait (wait_), speed (speed_ * 1000 / 60) {}

void GraphManager::add_stop(const string &stop) {
    stops.push_back (stop);
}

void GraphManager::relax_stops() {
    sort (stops.begin (), stops.end ());
    stops.resize (static_cast<unsigned long>(unique(stops.begin(), stops.end()) - stops.begin()));
//    dbg (stops);
    g = vector<vector<shared_ptr<edge>>> (stops.size());
}

int GraphManager::get_ind(const string &v) {
    return static_cast<int> (lower_bound (stops.begin (), stops.end (), v) - stops.begin ());
}

int GraphManager::add_edge(const string &from, const string &to, int len, const string& bus) {
    int from_ = get_ind (from), to_ = get_ind (to);
    shared_ptr<edge> e = make_shared<edge>(edge{to_, -1, static_cast<double>(len) / static_cast<double>(speed), static_cast<int>(edges.size()), bus});
    g[from_].push_back (e);
    edges.push_back (e);
//    cerr << "edge: " << from << ' ' << to << ' ' << len << ' ' << e->edge_num << '\n';
    return static_cast<int>(edges.size()) - 1;
}

void GraphManager::set_next(int edge_number, int next) {
    if (edge_number == -1) return;
    edges[edge_number]->next = next;
}

struct prev_info {
    int edge_num;
    int prev_vertex;
    double t;
    int cur_edge;
};

ostream& operator<< (ostream& out, const prev_info& x) {
    return out << x.edge_num << ' ' << x.prev_vertex << ' ' << x.t << ' ' << x.cur_edge;
}

double GraphManager::Dijkstra(const string &from, const string &to, ostream& out) {
    int s = get_ind (from), t = get_ind (to);
//    dbg (from);
//    dbg (s);
//    dbg (stops[s]);
    vector<unordered_map<int, double>> d (stops.size ());
    vector<unordered_map<int, vector<prev_info>>> prev (stops.size ());
    d[s][-1] = 0;
    set <pair<double, pair<int, int>>> a; // {distance, {vertex, last edge}}
    a.insert ({0, {s, -1}});
    while (!a.empty ()) {
        auto ptt = a.begin ()->second;
        double cur_dist = a.begin ()->first;
        a.erase (a.begin ());
        int v = ptt.first;
//        dbg (v);
//        dbg (stops[v]);
//        dbg (cur_dist);
//        dbg (ptt.second);
        if (ptt.second != -1) {
            int nxt = edges[ptt.second]->next;
            if (nxt != -1) {
                shared_ptr<edge> e = edges[nxt];
                if (!d[e->to].count (nxt) || d[e->to].at (nxt) > cur_dist + e->len) {
                    a.erase ({d[e->to][nxt], {e->to, nxt}});
                    d[e->to][nxt] = cur_dist + e->len;
                    a.insert ({d[e->to][nxt], {e->to, nxt}});
                    prev[e->to][nxt].clear ();
                    prev[e->to][nxt].push_back ({ptt.second, v, e->len, nxt});
                }
            }
        }
        for (auto& u: g[v]) {
            double new_dist = cur_dist + wait + u->len;
            int to_ = u->to;
            if (!d[to_].count (u->edge_num) || d[to_].at (u->edge_num) > new_dist) {
                a.erase ({d[to_][u->edge_num], {to_, u->edge_num}});
                d[to_][u->edge_num] = new_dist;
                a.insert ({d[to_][u->edge_num], {to_, u->edge_num}});
                prev[to_][u->edge_num].clear ();
                prev[to_][u->edge_num].push_back ({ptt.second, v, u->len, u->edge_num});
                prev[to_][u->edge_num].push_back ({-1, v, wait, -1});
            }
        }
    }
    double mn = 2e9;
    int edge_num = -1;
    for (auto& x: d[t]) {
        if (x.second < mn) {
            mn = x.second;
            edge_num = x.first;
        }
    }
//    cerr << from << ' ' << to << ' ' << mn << endl;
    if (mn == 2e9) {
        out << "\"error_message\": \"not found\"\n";
    } else {
        vector<prev_info> info;
        int last = t;
//        dbg (s);
        while (last != s) {
//            dbg (last);
//            dbg (edge_num);
            auto& v = prev[last][edge_num];
//            dbg (v.size ());
            for (auto& x: v) {
                info.push_back (x);
//                dbg (x);
            }
            edge_num = v.front ().edge_num;
            last = v.front ().prev_vertex;
        }
        reverse (info.begin (), info.end ());
//        dbg (info);
        out << "\"total_time\": " << mn << ",\n" << "\"items\": [\n";

        double time = 0;
        string bus;
        int cnt = 0;
        bool is_first = false;
        for (const auto& x: info) {
            if (x.cur_edge == -1) {
                if (time > 0) {
                    if (is_first) cout << ",\n";
                    out << "{\n" << "\"type\": \"Bus\",\n" << "\"bus\": " << '"' << bus << '"' << ",\n" <<
                    "\"span_count\": " << cnt << ",\n" << "\"time\": " << time << "\n}";
                }
                cnt = 0;
                time = 0;
                bus = "kek";
                if (is_first) cout << ",\n";
                is_first = true;
                out << "{\n" << "\"type\": \"Wait\",\n" << "\"stop_name\": " << '"' << stops[x.prev_vertex] << '"' << ",\n"
                << "\"time\": " << wait << "\n}";
            } else {
                ++cnt;
                bus = edges[x.cur_edge]->bus_num;
                time += x.t;
            }
        }
        if (is_first) cout << ",\n";
        if (time > 0) out << "{\n" << "\"type\": \"Bus\",\n" << "\"bus\": " << '"' << bus << '"' << ",\n" <<
            "\"span_count\": " << cnt << ",\n" << "\"time\": " << time << "\n}]\n";
        else out << "]\n";
    }
//    cerr << "DONE\n";
    return mn;
}