#include "prim.hpp"
#include "union_find.hpp"
#include <algorithm>
#include <vector>

Solution prim(const TSPLIBGraph& g, int k, int d, RNG& rng) {
    const int n = g.size();
    Solution sol(g);

    if (n <= 1 || k <= 0) return sol;

    std::vector<bool> in_tree(n, false);

    // Start from random node
    in_tree[static_cast<uint32_t>(rng.next_int(n))] = true;
    // Collect candidate edges: (tree node) -- (non-tree node), both degree < d
    std::vector<size_t> candidates;

    for (int added = 0; added < k - 1; ++added) {
        candidates.clear();

        for (int u = 0; u < n; ++u) {
            if (!in_tree[u]) continue;
            if (sol.degree[u] >= d) continue;
            for (int v = 0; v < n; ++v) {
                if (in_tree[v]) continue;
                if (sol.degree[v] >= d) continue;
                candidates.push_back(g.edge_id(u, v));
            }
        }

        if (candidates.empty()) break;

        // Pick a random candidate
        size_t eid = candidates[rng.next_int(static_cast<uint32_t>(candidates.size()))];
        auto [a, b] = g.edge_nodes(eid);

        // Mark the non-tree endpoint as in-tree
        if (!in_tree[a]) in_tree[a] = true;
        else             in_tree[b] = true;

        sol.add_edge(eid);
    }

    return sol;
}

Solution prim_original(const TSPLIBGraph& g, int k, int d, RNG& rng) {
    const int n = g.size();

    // ── Step 1: Prim's MST (unrestricted) ──────────────────────────────────
    std::vector<bool> in_mst(n, false);
    std::vector<double> min_cost(n, 1e300);
    std::vector<int> min_from(n, -1);

    int start = static_cast<int>(rng.next_int(static_cast<uint32_t>(n)));
    min_cost[start] = 0.0;

    Solution mst(g);

    for (int added = 0; added < n; ++added) {
        int u = -1;
        for (int v = 0; v < n; ++v) {
            if (!in_mst[v] && (u < 0 || min_cost[v] < min_cost[u]))
                u = v;
        }
        in_mst[u] = true;
        if (min_from[u] >= 0)
            mst.add_edge(g.edge_id(u, min_from[u]));

        for (int v = 0; v < n; ++v) {
            if (!in_mst[v] && g.dist(u, v) < min_cost[v]) {
                min_cost[v] = g.dist(u, v);
                min_from[v] = u;
            }
        }
    }

    // ── Step 2: remove edges violating degree constraint ───────────────────
    std::vector<size_t> edges_snapshot(mst.edges.begin(), mst.edges.end());
    for (int i = static_cast<int>(edges_snapshot.size()) - 1; i > 0; --i) {
        int j = static_cast<int>(rng.next_int(static_cast<uint32_t>(i + 1)));
        std::swap(edges_snapshot[i], edges_snapshot[j]);
    }
    for (size_t e : edges_snapshot) {
        auto [u, v] = g.edge_nodes(e);
        if (mst.degree[u] > d || mst.degree[v] > d)
            mst.remove_edge(e);
    }

    // ── Step 3: Union-Find on remaining edges ──────────────────────────────
    UnionFind uf(n);
    for (size_t e : mst.edges) {
        auto [u, v] = g.edge_nodes(e);
        uf.unite(u, v);
    }

    const auto& sorted_edges = g.sorted_edges();

    // 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;
        }
    }

    // Greedily add edges connecting different components, respecting d
    for (size_t e : sorted_edges) {
        if (biggest >= k) break;
        if (mst.has_edge(e)) continue;
        auto [u, v] = g.edge_nodes(e);
        if (mst.degree[u] >= d || mst.degree[v] >= d) continue;
        if (uf.find(u) == uf.find(v)) continue;
        mst.add_edge(e);
        int root = uf.unite(u, v);
        if (uf.sz[root] > biggest) {
            biggest = uf.sz[root];
            best_root = root;
        }
    }

    // ── Step 4: keep only the largest component ────────────────────────────
    for (int i = static_cast<int>(mst.edges.size()) - 1; i >= 0; --i) {
        size_t e = mst.edges[i];
        auto [u, v] = g.edge_nodes(e);
        if (uf.find(u) != best_root)
            mst.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 (mst.degree[v] != 1) continue;
            size_t e = g.edge_id(v, mst.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;
        mst.remove_edge(worst_edge);
        --node_count;
    }

    return mst;
}