#include "vnd1.hpp"
#include <vector>

namespace {

struct TreeInfo {
    std::vector<int> par;
    std::vector<int> tin, tout;

    void build(const Solution& sol) {
        int n = sol.graph->size();
        par.assign(n, -1);
        tin.assign(n, -1);
        tout.assign(n, -1);

        int root = -1;
        for (int v = 0; v < n; ++v)
            if (sol.degree[v] > 0) { root = v; break; }
        if (root < 0) return;

        struct Frame { int node; int idx; };
        std::vector<Frame> stack;
        int timer = 0;
        tin[root] = timer++;
        stack.push_back({root, 0});

        while (!stack.empty()) {
            auto& [node, idx] = stack.back();
            if (idx < static_cast<int>(sol.adj[node].size())) {
                int child = sol.adj[node][idx++];
                if (tin[child] < 0) {
                    par[child] = node;
                    tin[child] = timer++;
                    stack.push_back({child, 0});
                }
            } else {
                tout[node] = timer++;
                stack.pop_back();
            }
        }
    }

    bool in_subtree(int ancestor, int node) const {
        return tin[ancestor] <= tin[node] && tout[node] <= tout[ancestor];
    }

    bool swap_connected(int u1, int u2, int v1, int v2) const {
        int c1 = (par[u2] == u1) ? u2 : u1;
        int c2 = (par[v2] == v1) ? v2 : v1;

        auto comp_of = [&](int node) -> int {
            return (in_subtree(c1, node) ? 1 : 0)
                 | (in_subtree(c2, node) ? 2 : 0);
        };

        int cu1 = comp_of(u1), cv1 = comp_of(v1);
        int cu2 = comp_of(u2), cv2 = comp_of(v2);

        if (cu1 == cv1 || cu2 == cv2) return false;

        bool seen[4] = {};
        seen[cu1] = seen[cv1] = seen[cu2] = seen[cv2] = true;
        return (seen[0] + seen[1] + seen[2] + seen[3]) >= 3;
    }
};

} // anonymous namespace

bool vnd1(Solution& sol) {
    const auto& g = *sol.graph;
    const int n = g.size();

    TreeInfo tree;
    tree.build(sol);

    for (int u1 = 0; u1 < n; ++u1) {
        for (int v1 = u1 + 1; v1 < n; ++v1) {
            for (int u2 : sol.adj[u1]) {
                if (u2 == v1) continue;
                for (int v2 : sol.adj[v1]) {
                    if (v2 == u1 || v2 == u2) continue;

                    if (sol.has_edge(g.edge_id(u1, v1))) continue;
                    if (sol.has_edge(g.edge_id(u2, v2))) continue;

                    double old_cost = g.dist(u1, u2) + g.dist(v1, v2);
                    double new_cost = g.dist(u1, v1) + g.dist(u2, v2);

                    if (new_cost < old_cost) {
                        if (!tree.swap_connected(u1, u2, v1, v2)) continue;
                        sol.remove_edge(g.edge_id(u1, u2));
                        sol.remove_edge(g.edge_id(v1, v2));
                        sol.add_edge(g.edge_id(u1, v1));
                        sol.add_edge(g.edge_id(u2, v2));
                        return true;
                    }
                }
            }
        }
    }

    return false;
}

bool vnd1bi(Solution& sol) {
    const auto& g = *sol.graph;
    const int n = g.size();

    TreeInfo tree;
    tree.build(sol);

    double best_gain = 0.0;
    int best_u1 = -1, best_v1 = -1, best_u2 = -1, best_v2 = -1;

    for (int u1 = 0; u1 < n; ++u1) {
        for (int v1 = u1 + 1; v1 < n; ++v1) {
            for (int u2 : sol.adj[u1]) {
                if (u2 == v1) continue;
                for (int v2 : sol.adj[v1]) {
                    if (v2 == u1 || v2 == u2) continue;

                    if (sol.has_edge(g.edge_id(u1, v1))) continue;
                    if (sol.has_edge(g.edge_id(u2, v2))) continue;

                    double gain = g.dist(u1, u2) + g.dist(v1, v2)
                                - g.dist(u1, v1) - g.dist(u2, v2);

                    if (gain > best_gain) {
                        if (!tree.swap_connected(u1, u2, v1, v2)) continue;
                        best_gain = gain;
                        best_u1 = u1; best_v1 = v1;
                        best_u2 = u2; best_v2 = v2;
                    }
                }
            }
        }
    }

    if (best_u1 < 0) return false;

    sol.remove_edge(g.edge_id(best_u1, best_u2));
    sol.remove_edge(g.edge_id(best_v1, best_v2));
    sol.add_edge(g.edge_id(best_u1, best_v1));
    sol.add_edge(g.edge_id(best_u2, best_v2));
    return true;
}