#include "lib/exec.h"

#include <iterator>
#include <ranges>
#include <tuple>
#include <vector>

using namespace std;

template<typename T, typename Iter>
    requires random_access_iterator<Iter> && is_same_v<remove_cv_t<iter_value_t<Iter>>, remove_cv_t<T>>
class MDSpan {
    const Iter begin, end;

    public:

    const size_t Nx, Ny;

    template<typename R>
        requires is_same_v<ranges::iterator_t<R>, Iter>
    MDSpan(R&& rng, size_t Nx, size_t Ny) : begin(std::begin(rng)), end(std::end(rng)), Nx(Nx), Ny(Ny) {
        if (static_cast<signed long>(Nx * Ny) != end - begin) {
            throw string("bad dimensions");
        }
    }

    MDSpan(const MDSpan<T, Iter>& other) : begin(other.begin), end(other.end), Nx(other.Nx), Ny(other.Ny) {
        fmt::print("copy ctor\n");
    }

    const MDSpan<T, Iter>& operator=(const MDSpan<T, Iter>& other) = delete;

    T& operator[](size_t x, size_t y) {
        size_t ix = y * Nx + x;
        return begin[ix];
    }

    void print() {
        for (size_t y = 0; y < Ny; y++) {
            for (size_t x = 0; x < Nx; x++) {
                fmt::print("{:5} ", (*this)[x, y]);
            }
            fmt::print("\n");
        }
    }
};

template<typename Cont> requires ranges::range<Cont>
MDSpan<ranges::range_value_t<Cont>, ranges::iterator_t<Cont>> make_mdspan(Cont&& c, size_t x, size_t y) {
    return MDSpan<ranges::range_value_t<Cont>, ranges::iterator_t<Cont>>(c, x, y);
}

template<typename C>
concept constant_range = requires (C c) {
    { c.cbegin() } -> same_as<typename remove_cvref_t<C>::const_iterator>;
    { c.cend() } -> same_as<typename remove_cvref_t<C>::const_iterator>;
};

// Note: const specialization, uses static member type
template <typename Cont>
  requires constant_range<Cont>
MDSpan<const ranges::range_value_t<remove_cvref_t<Cont>>,
       typename remove_cvref_t<Cont>::const_iterator>
make_const_mdspan(Cont &&c, size_t x, size_t y) {
    return MDSpan<const ranges::range_value_t<remove_cvref_t<Cont>>,
                  typename remove_cvref_t<Cont>::const_iterator>(c, x, y);
}

// https://www.techiedelight.com/single-source-shortest-paths-bellman-ford-algorithm/
void bellman_ford() {
    const size_t nodes = 5;
    const int inf = numeric_limits<int>::max();
    vector<int> v(nodes * nodes);

    //MDSpan<const int, vector<int>::const_iterator> mds = make_mdspan(v, 6, 7);
    auto adj = make_mdspan(v, nodes, nodes);

    adj[0, 1] = -1;
    adj[0, 2] = 4;
    adj[1, 2] = 3;
    adj[1, 3] = 2;
    adj[3, 1] = 1;
    adj[1, 4] = 2;
    adj[4, 3] = -3;
    adj[3, 2] = 5;

    adj.print();
    vector<int> cost(nodes, inf);
    cost[0] = 0;

    for (int i = 0; i < nodes-1; i++) {
        for (size_t a = 0; a < nodes; a++) {
            for (size_t b = 0; b < nodes; b++) {
                // Not elegant. Need separate store for adjacency, basically.
                if (a == b || adj[a, b] == 0) continue;

                if (cost[a] + adj[a,b] < cost[b]) cost[b] = cost[a] + adj[a, b];
            }
        }
    }

    fmt::print("{}\n", cost);
}

void floyd_warshall() {
    const size_t nodes = 4;
    const int inf = numeric_limits<int>::max();
    vector<tuple<int, int, int>> edges{
        {1, 0, 4},
        {1, 2, 3},
        {3, 1, -1},
        {0, 2, -2},
        {2, 3, 2},
    };

    vector<int> adj_v(nodes*nodes);
    auto adj = make_mdspan(adj_v, nodes, nodes);

    for (const auto& e : edges){
        adj[get<0>(e), get<1>(e)] = get<2>(e);
    }
    for (size_t i = 0; i < nodes; i++) {
        for (size_t j = 0; j < nodes; j++) {
            if (i == j) adj[i,j] = 0;
            else if (adj[i,j] == 0) adj[i,j] = inf;
        }
    }

    adj.print();

    for (size_t j = 0; j < nodes; j++) {
        for (size_t i = 0; i < nodes; i++) {
            for (size_t k = 0; k < nodes; k++) {
                if (adj[i, j] != inf && adj[j, k] != inf) {
                    if (adj[i,j] + adj[j,k] < adj[i,k]) adj[i,k] = adj[i,j]+adj[j,k];
                }
            }
        }
    }

    fmt::print("\n");
    adj.print();
}

int main(void) {
    timeit_f([]() { bellman_ford(); }, "bellman-ford");
    floyd_warshall();
    return 0;
}