#include "shake.hpp"
#include "union_find.hpp"
#include <algorithm>
#include <vector>
#include <cmath>

// ── shake2: destructive rebuild ─────────────────────────────────────────────
//
// 1. Remove k_shake random edges from the current solution.
// 2. Treat remaining edges + isolated nodes as connected components (Union-Find).
// 3. Add random edges connecting different components (respecting degree d)
//    until one component has >= k nodes.
// 4. Keep only the largest component, remove all other edges.
// 5. Prune random leaves until exactly k nodes.
//
// Time complexity:  O(n · C + n · (comp_size - k))
//   where C = number of components after step 1, comp_size = largest component.
//
// Step 1 (remove k_shake edges):       O(k_shake)  (random access + swap-and-pop)
// Step 2 (build UF):                   O(k · α(n)) (iterate remaining edges)
// Step 3 (candidate grouping):         O(n)        (one pass over all nodes)
//        (biggest component scan):     O(n)        (one pass)
//        (merge loop, per iteration):  O(n)        (active_roots rebuild)
//          × at most C-1 merges:       O(n · C)    worst case
//          candidate list merges:      O(n) total  (smaller-into-larger)
// Step 4 (keep largest component):     O(k)        (backward scan over sol.edges)
// Step 5 (prune leaves):               O(n) per pass, at most O(k) passes
//          but each pass removes ≥1:   O(n · (comp_size - k)) worst case
//
// Total: O(n · C + n · (comp_size - k))
//   In practice C is small (≤ k_shake + 1) and comp_size - k is small,
//   so the dominant cost is the O(n) scans in steps 3 and 5.

void shake2(Solution& sol, int k_shake, int k, int d, RNG& rng) {
    const auto& g = *sol.graph;
    const int n = g.size();

    // ── Step 1: remove k_shake random edges ────────────────────────────────
    int to_remove = std::min(k_shake, static_cast<int>(sol.edges.size()));
    for (int i = 0; i < to_remove; ++i) {
        int idx = static_cast<int>(rng.next_int(
            static_cast<uint32_t>(sol.edges.size())));
        sol.remove_edge(sol.edges[idx]);
    }

    // ── Step 2: build UF from remaining edges ──────────────────────────────
    UnionFind uf(n);
    for (size_t e : sol.edges) {
        auto [u, v] = g.edge_nodes(e);
        uf.unite(u, v);
    }

    // ── Step 3: add random cross-component edges until one comp >= k ───────
    //
    // Group candidate nodes (degree < d) by component root.
    // To merge two components: pick two different groups, pick one node from
    // each — guaranteed cross-component edge, no wasted attempts.

    // comp_cands[root] = shuffled list of candidate nodes in that component.
    // active_roots = list of roots that have at least one candidate node.
    std::vector<std::vector<int>> comp_cands(n);
    std::vector<int> active_roots;

    for (int v = 0; v < n; ++v) {
        if (sol.degree[v] < d)
            comp_cands[uf.find(v)].push_back(v);
    }
    for (int r = 0; r < n; ++r) {
        if (!comp_cands[r].empty()) active_roots.push_back(r);
    }

    // Find current biggest component.
    int biggest = -1;
    int best_root = 0;
    for (int v = 0; v < n; ++v) {
        int root = uf.find(v);
        if (uf.sz[root] > biggest) {
            biggest = uf.sz[root];
            best_root = root;
        }
    }

    while (biggest < k && active_roots.size() >= 2) {
        // Pick two distinct active roots.
        int ri = static_cast<int>(rng.next_int(
            static_cast<uint32_t>(active_roots.size())));
        int rj = static_cast<int>(rng.next_int(
            static_cast<uint32_t>(active_roots.size() - 1)));
        if (rj >= ri) ++rj;

        int root_a = active_roots[ri], root_b = active_roots[rj];

        // Pick random candidate from each.
        auto& ca = comp_cands[root_a];
        auto& cb = comp_cands[root_b];
        int u = ca[rng.next_int(static_cast<uint32_t>(ca.size()))];
        int v = cb[rng.next_int(static_cast<uint32_t>(cb.size()))];

        sol.add_edge(g.edge_id(u, v));
        int merged_root = uf.unite(u, v);

        // Determine which root survived and which was absorbed.
        int surv = merged_root;
        // Find which original root maps to survivor vs absorbed.
        int orig_surv, orig_dead;
        if (uf.find(root_a) == surv) {
            orig_surv = root_a; orig_dead = root_b;
        } else {
            orig_surv = root_b; orig_dead = root_a;
        }

        // Merge candidate lists into the survivor's slot.
        // Move the smaller into the larger.
        auto& cs = comp_cands[orig_surv];
        auto& cd = comp_cands[orig_dead];
        if (cs.size() < cd.size()) {
            cd.insert(cd.end(), cs.begin(), cs.end());
            cs.clear();
            std::swap(comp_cands[surv], cd);  // put merged list at surv
            if (surv != orig_dead) comp_cands[orig_dead].clear();
        } else {
            cs.insert(cs.end(), cd.begin(), cd.end());
            cd.clear();
            if (surv != orig_surv) {
                std::swap(comp_cands[surv], cs);
                comp_cands[orig_surv].clear();
            }
        }

        // Remove u/v from candidate list if they hit degree d.
        auto remove_from = [](std::vector<int>& vec, int val) {
            for (int i = static_cast<int>(vec.size()) - 1; i >= 0; --i) {
                if (vec[i] == val) {
                    vec[i] = vec.back(); vec.pop_back(); return;
                }
            }
        };
        if (sol.degree[u] >= d) remove_from(comp_cands[surv], u);
        if (sol.degree[v] >= d) remove_from(comp_cands[surv], v);

        if (uf.sz[surv] > biggest) {
            biggest = uf.sz[surv];
            best_root = surv;
        }

        // Rebuild active_roots: remove dead, ensure surv is present.
        // Just rebuild — list is small.
        active_roots.clear();
        for (int r = 0; r < n; ++r) {
            if (!comp_cands[r].empty() && uf.find(r) == r)
                active_roots.push_back(r);
        }
    }

    // ── Step 4: keep only the largest component ────────────────────────────
    for (int i = static_cast<int>(sol.edges.size()) - 1; i >= 0; --i) {
        size_t e = sol.edges[i];
        auto [u, v] = g.edge_nodes(e);
        if (uf.find(u) != best_root)
            sol.remove_edge(e);
    }

    // ── Step 5: prune leaves until exactly k nodes ─────────────────────────
    int node_count = uf.sz[best_root];

    while (node_count > k) {
        std::vector<int> leaves;
        for (int v = 0; v < n; ++v) {
            if (sol.degree[v] == 1) leaves.push_back(v);
        }
        if (leaves.empty()) break;

        // Shuffle for random pruning order.
        for (int i = static_cast<int>(leaves.size()) - 1; i > 0; --i) {
            int j = static_cast<int>(rng.next_int(static_cast<uint32_t>(i + 1)));
            std::swap(leaves[i], leaves[j]);
        }

        for (int v : leaves) {
            if (node_count <= k) break;
            if (sol.degree[v] != 1) continue;
            int nb = sol.adj[v][0];
            sol.remove_edge(g.edge_id(v, nb));
            --node_count;
        }
    }
}

// ── shake3: GRASP-style destructive rebuild ─────────────────────────────────
//
// Like shake2, but step 3 uses a greedy-randomized (GRASP) approach:
// iterate through sorted graph edges, skip each eligible edge with
// probability 0.15. Repeat passes over skipped edges until one
// component reaches k nodes.
//
// Time complexity: O(n^2 log n + n · (comp_size - k))
//   Dominated by sorting all graph edges: O(n^2 log n).
//   Each merge pass is O(|candidates|) where candidates shrink each pass.
//   Steps 4 and 5 same as shake2.

void shake3(Solution& sol, int k_shake, int k, int d, RNG& rng) {
    const auto& g = *sol.graph;
    const int n = g.size();

    // ── Step 1: remove k_shake random edges ────────────────────────────────
    int to_remove = std::min(k_shake, static_cast<int>(sol.edges.size()));
    for (int i = 0; i < to_remove; ++i) {
        int idx = static_cast<int>(rng.next_int(
            static_cast<uint32_t>(sol.edges.size())));
        sol.remove_edge(sol.edges[idx]);
    }

    // ── Step 2: build UF from remaining edges ──────────────────────────────
    UnionFind uf(n);
    for (size_t e : sol.edges) {
        auto [u, v] = g.edge_nodes(e);
        uf.unite(u, v);
    }

    // Find current biggest component.
    int biggest = -1;
    int best_root = 0;
    for (int v = 0; v < n; ++v) {
        int root = uf.find(v);
        if (uf.sz[root] > biggest) {
            biggest = uf.sz[root];
            best_root = root;
        }
    }

    // ── Step 3: GRASP edge addition ────────────────────────────────────────
    std::vector<size_t> candidates(g.sorted_edges().begin(), g.sorted_edges().end());
//    unsigned int alpha = 100 * std::min(0.85, std::sqrt(3.5 * k_shake * k) / std::sqrt(100.0 * d * n));
    unsigned int alpha = 15;

    while (biggest < k && !candidates.empty()) {
        std::vector<size_t> skipped;
        for (size_t e : candidates) {
            if (biggest >= k) break;
            auto [u, v] = g.edge_nodes(e);
            if (uf.find(u) == uf.find(v)) continue;
            if (sol.degree[u] >= d || sol.degree[v] >= d) continue;
            if (sol.has_edge(e)) continue;
            // GRASP: skip with probability 0.15
            if (rng.next_int(100) < alpha) {
                skipped.push_back(e);
                continue;
            }
            sol.add_edge(e);
            int root = uf.unite(u, v);
            if (uf.sz[root] > biggest) {
                biggest = uf.sz[root];
                best_root = root;
            }
        }
        candidates = std::move(skipped);
    }

    // ── Step 4: keep only the largest component ────────────────────────────
    for (int i = static_cast<int>(sol.edges.size()) - 1; i >= 0; --i) {
        size_t e = sol.edges[i];
        auto [u, v] = g.edge_nodes(e);
        if (uf.find(u) != best_root)
            sol.remove_edge(e);
    }

    // ── Step 5: prune longest leaf edges until exactly k nodes ─────────────
    int node_count = uf.sz[best_root];

    while (node_count > k) {
        int worst_leaf = -1;
        size_t worst_edge = 0;
        double worst_cost = -1.0;
        for (int v = 0; v < n; ++v) {
            if (sol.degree[v] != 1) continue;
            size_t e = g.edge_id(v, sol.adj[v][0]);
            if (g.dist(e) > worst_cost) {
                worst_cost = g.dist(e);
                worst_leaf = v;
                worst_edge = e;
            }
        }
        if (worst_leaf < 0) break;
        sol.remove_edge(worst_edge);
        --node_count;
    }
}

// Time complexity: O(n + vns_k)
//
// Initialization (deg0/deg1 build):  O(n)
// Per iteration:                     O(1) amortized
//   - random edge from sol.edges:    O(1)  (vector random access)
//   - random leaf from deg1:         O(1)  (rejection: at most 2 bad nodes)
//   - leaf's edge via adj:           O(1)  (leaf has exactly 1 neighbor)
//   - random deg0 node:              O(1)  (indexed list random access)
//   - all deg0/deg1 updates:         O(1)  (indexed list add/remove)
//   - sol add/remove_edge:           O(1)  amortized
//
// Total: O(n + vns_k)

void shake(Solution& sol, int vns_k, RNG& rng) {
    const auto& g = *sol.graph;
    const int n = g.size();

    // ── Indexed list: O(1) add / remove / random access ─────────────────
    //    items vector + pos array; swap-and-pop for removal
    struct IndexedList {
        std::vector<int> items;
        std::vector<int> pos;   // pos[v] = index in items, or -1

        explicit IndexedList(int n) : pos(n, -1) {}

        void add(int v) {
            if (pos[v] != -1) return;
            pos[v] = static_cast<int>(items.size());
            items.push_back(v);
        }
        void remove(int v) {
            if (pos[v] == -1) return;
            int i = pos[v];
            int last = items.back();
            items[i] = last;
            pos[last] = i;
            items.pop_back();
            pos[v] = -1;
        }
        bool contains(int v) const { return pos[v] != -1; }
        bool empty()         const { return items.empty(); }
        int  size()          const { return static_cast<int>(items.size()); }
        int  random(RNG& r)  const { return items[r.next_int(static_cast<uint32_t>(items.size()))]; }
    };

    IndexedList deg0(n), deg1(n);

    for (int v = 0; v < n; ++v) {
        if      (sol.degree[v] == 0) deg0.add(v);
        else if (sol.degree[v] == 1) deg1.add(v);
    }

    // ── Degree-transition helper ────────────────────────────────────────
    //    Called after sol.add_edge / sol.remove_edge has already updated
    //    sol.degree.  old_deg is the degree before that operation.
    auto update = [&](int v, int old_deg) {
        int new_deg = sol.degree[v];
        if (old_deg == new_deg) return; // no change
        if (old_deg == 0)      deg0.remove(v);
        else if (old_deg == 1) deg1.remove(v);
        if (new_deg == 0)      deg0.add(v);
        else if (new_deg == 1) deg1.add(v);
    };

    // Wrappers that track old degrees and update indexed lists
    auto remove_edge = [&](size_t e) {
        auto [a, b] = g.edge_nodes(e);
        int da = sol.degree[a], db = sol.degree[b];
        sol.remove_edge(e);
        update(a, da);
        update(b, db);
    };

    auto add_edge = [&](size_t e) {
        auto [a, b] = g.edge_nodes(e);
        int da = sol.degree[a], db = sol.degree[b];
        sol.add_edge(e);
        update(a, da);
        update(b, db);
    };

    // ── Main loop ───────────────────────────────────────────────────────
    for (int iter = 0; iter < vns_k; ++iter) {
        if (sol.edges.empty()) break;

        // 1. Pick a random edge
        size_t rand_edge = sol.edges[rng.next_int(static_cast<uint32_t>(sol.edges.size()))];
        auto [eu, ev] = g.edge_nodes(rand_edge);

        // 2. Pick a random leaf not on that edge (rejection sampling, O(1) expected)
        int bad = (sol.degree[eu] == 1 ? 1 : 0) + (sol.degree[ev] == 1 ? 1 : 0);
        if (deg1.size() <= bad) continue;

        int leaf;
        do {
            leaf = deg1.random(rng);
        } while (leaf == eu || leaf == ev);

        // 3. Remove the leaf's edge (O(1) via adjacency list)
        int leaf_neighbor = sol.adj[leaf][0];
        size_t leaf_edge = g.edge_id(leaf, leaf_neighbor);
        remove_edge(leaf_edge);

        // 4. Pick a random degree-0 node
        if (deg0.empty()) continue;
        int new_node = deg0.random(rng);

        // 5. Split rand_edge through new_node
        remove_edge(rand_edge);
        add_edge(g.edge_id(eu, new_node));
        add_edge(g.edge_id(ev, new_node));
    }
}