#pragma once
#include "tsplib_graph.hpp"
#include <queue>
#include <unordered_map>
#include <vector>

struct Solution {
    const TSPLIBGraph*                  graph;     // non-owning, never null
    std::vector<size_t>                 edges;     // edge ids, random-accessible
    std::unordered_map<size_t, size_t>  edge_pos;  // edge_id -> index in edges
    std::vector<std::vector<int>>       adj;       // adj[v] = neighbors in solution
    std::vector<int>                    degree;    // degree[v] = adj[v].size()
    double                              cost;

    explicit Solution(const TSPLIBGraph& g)
        : graph(&g), adj(g.size()), degree(g.size(), 0), cost(0.0) {}

    void add_edge(size_t e) {
        edge_pos[e] = edges.size();
        edges.push_back(e);
        cost += graph->dist(e);
        auto [u, v] = graph->edge_nodes(e);
        degree[u]++;
        degree[v]++;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }

    void remove_edge(size_t e) {
        // Swap-and-pop from edges vector
        auto it = edge_pos.find(e);
        size_t idx = it->second;
        size_t last = edges.back();
        edges[idx] = last;
        edge_pos[last] = idx;
        edges.pop_back();
        edge_pos.erase(it);

        cost -= graph->dist(e);
        auto [u, v] = graph->edge_nodes(e);
        degree[u]--;
        degree[v]--;

        // Remove v from adj[u] (swap-and-pop)
        auto& au = adj[u];
        for (size_t i = 0; i < au.size(); ++i) {
            if (au[i] == v) { au[i] = au.back(); au.pop_back(); break; }
        }
        // Remove u from adj[v]
        auto& av = adj[v];
        for (size_t i = 0; i < av.size(); ++i) {
            if (av[i] == u) { av[i] = av.back(); av.pop_back(); break; }
        }
    }

    bool has_edge(size_t e) const { return edge_pos.count(e); }

    // ── rooted tree representation ──────────────────────────────────────────
    std::vector<int> parent;  // parent[v] in rooted tree, -1 if not in tree
    std::vector<int> depth;   // depth[v] from root, -1 if not in tree

    void build_rooted_tree() {
        int n = graph->size();
        parent.assign(n, -1);
        depth.assign(n, -1);
        int root = -1;
        for (int v = 0; v < n; ++v)
            if (degree[v] > 0) { root = v; break; }
        if (root < 0) return;
        depth[root] = 0;
        std::queue<int> q;
        q.push(root);
        while (!q.empty()) {
            int u = q.front(); q.pop();
            for (int w : adj[u]) {
                if (depth[w] < 0) {
                    parent[w] = u;
                    depth[w] = depth[u] + 1;
                    q.push(w);
                }
            }
        }
    }
};