// Grafo de estados: lo que hay que recordar entra en el nodo.
// Aquí se recuerda cuántos tramos lleva el recorrido, módulo c.
#pragma once

#include <cstddef>
#include <functional>
#include <queue>
#include <utility>
#include <vector>

#include "shortest_paths.h"

struct WalkResult {
    long long distance;      // UNREACHABLE si no hay recorrido
    std::vector<int> walk;   // las ciudades del recorrido, con repeticiones
    long long pops;          // entradas sacadas del heap
    long long relaxations;   // arcos del grafo de estados que se miraron
};

// Recorrido mínimo de source a target con un número de tramos múltiplo de c.
// El estado (ciudad, resto) se guarda como ciudad * c + resto. Los arcos no se construyen:
// salen de la lista de la ciudad, y todos van del resto r al resto (r + 1) mod c.
inline WalkResult min_walk_multiple(const Graph& graph, int source, int target, int c) {
    using Entry = std::pair<long long, std::size_t>;
    std::size_t layers = static_cast<std::size_t>(c);
    std::size_t states = graph.size() * layers;
    std::vector<long long> distance(states, UNREACHABLE);
    std::vector<std::size_t> parent(states, states);
    std::vector<bool> fixed(states, false);
    std::priority_queue<Entry, std::vector<Entry>, std::greater<Entry>> heap;
    std::size_t start = static_cast<std::size_t>(source) * layers;
    distance[start] = 0;
    heap.push({0, start});
    WalkResult result{UNREACHABLE, {}, 0, 0};
    while (!heap.empty()) {
        auto [d, state] = heap.top();
        heap.pop();
        result.pops++;
        if (fixed[state]) continue;
        fixed[state] = true;
        std::size_t city = state / layers;
        std::size_t next_rest = (state % layers + 1) % layers;
        for (const Arc& arc : graph[city]) {
            result.relaxations++;
            std::size_t next = static_cast<std::size_t>(arc.to) * layers + next_rest;
            if (d + arc.weight < distance[next]) {
                distance[next] = d + arc.weight;
                parent[next] = state;
                heap.push({distance[next], next});
            }
        }
    }
    std::size_t goal = static_cast<std::size_t>(target) * layers;
    result.distance = distance[goal];
    if (distance[goal] != UNREACHABLE) {
        std::vector<int> reversed;
        std::size_t state = goal;
        while (state != start) {
            reversed.push_back(static_cast<int>(state / layers));
            state = parent[state];
        }
        reversed.push_back(source);
        result.walk.assign(reversed.rbegin(), reversed.rend());
    }
    return result;
}