#include "vnd2.hpp"
#include <algorithm>
#include <vector>

// Walk both nodes up until they meet.
static int lca(const Solution& sol, int u, int v) {
    while (u != v) {
        if (sol.depth[u] > sol.depth[v])
            u = sol.parent[u];
        else
            v = sol.parent[v];
    }
    return u;
}

// Walk target up until parent[w] == ancestor, return w.
static int child_toward(const Solution& sol, int ancestor, int target) {
    int w = target;
    while (sol.parent[w] != ancestor)
        w = sol.parent[w];
    return w;
}

// Walk u→ancestor and v→ancestor, return edge_id of the max-weight edge.
static size_t max_edge_on_path(const TSPLIBGraph& g, const Solution& sol,
                               int u, int v, int ancestor) {
    size_t best_eid = 0;
    double best_cost = -1.0;

    int w = u;
    while (w != ancestor) {
        size_t eid = g.edge_id(w, sol.parent[w]);
        double c = g.dist(eid);
        if (c > best_cost) { best_cost = c; best_eid = eid; }
        w = sol.parent[w];
    }
    w = v;
    while (w != ancestor) {
        size_t eid = g.edge_id(w, sol.parent[w]);
        double c = g.dist(eid);
        if (c > best_cost) { best_cost = c; best_eid = eid; }
        w = sol.parent[w];
    }
    return best_eid;
}

bool vnd2(Solution& sol, int d, RNG& rng) {
    const auto& g = *sol.graph;
    const int n = g.size();

    sol.build_rooted_tree();

    std::vector<int> perm(n);
    for (int i = 0; i < n; ++i) perm[i] = i;

    // Fisher-Yates shuffle
    for (int i = n - 1; i > 0; --i) {
        int j = static_cast<int>(rng.next_int(static_cast<uint32_t>(i + 1)));
        std::swap(perm[i], perm[j]);
    }

    for (int pi = 0; pi < n; ++pi) {
        int u = perm[pi];
        if (sol.depth[u] < 0) continue;

        for (int qi = pi + 1; qi < n; ++qi) {
            int v = perm[qi];
            if (sol.depth[v] < 0) continue;

            size_t add_eid = g.edge_id(u, v);
            if (sol.has_edge(add_eid)) continue;
            if (sol.degree[u] >= d && sol.degree[v] >= d) continue;

            int ancestor = lca(sol, u, v);

            size_t remove_eid;
            if (sol.degree[u] >= d) {
                // Must remove cycle edge incident to u
                int other = (u == ancestor) ? child_toward(sol, u, v)
                                            : sol.parent[u];
                remove_eid = g.edge_id(u, other);
            } else if (sol.degree[v] >= d) {
                // Must remove cycle edge incident to v
                int other = (v == ancestor) ? child_toward(sol, v, u)
                                            : sol.parent[v];
                remove_eid = g.edge_id(v, other);
            } else {
                // Remove max-weight edge on the cycle
                remove_eid = max_edge_on_path(g, sol, u, v, ancestor);
            }

            if (g.dist(add_eid) < g.dist(remove_eid)) {
                sol.remove_edge(remove_eid);
                sol.add_edge(add_eid);
                return true;
            }
        }
    }

    return false;
}

bool vnd2bi(Solution& sol, int d, [[maybe_unused]] RNG& rng) {
    const auto& g = *sol.graph;
    const int n = g.size();

    sol.build_rooted_tree();

    double best_gain = 0.0;
    size_t best_add = 0, best_remove = 0;

    for (int u = 0; u < n; ++u) {
        if (sol.depth[u] < 0) continue;

        for (int v = u + 1; v < n; ++v) {
            if (sol.depth[v] < 0) continue;

            size_t add_eid = g.edge_id(u, v);
            if (sol.has_edge(add_eid)) continue;
            if (sol.degree[u] >= d && sol.degree[v] >= d) continue;

            int ancestor = lca(sol, u, v);

            size_t remove_eid;
            if (sol.degree[u] >= d) {
                int other = (u == ancestor) ? child_toward(sol, u, v)
                                            : sol.parent[u];
                remove_eid = g.edge_id(u, other);
            } else if (sol.degree[v] >= d) {
                int other = (v == ancestor) ? child_toward(sol, v, u)
                                            : sol.parent[v];
                remove_eid = g.edge_id(v, other);
            } else {
                remove_eid = max_edge_on_path(g, sol, u, v, ancestor);
            }

            double gain = g.dist(remove_eid) - g.dist(add_eid);
            if (gain > best_gain) {
                best_gain = gain;
                best_add = add_eid;
                best_remove = remove_eid;
            }
        }
    }

    if (best_gain <= 0.0) return false;

    sol.remove_edge(best_remove);
    sol.add_edge(best_add);
    return true;
}