#include <c_BoruvkaMST.h>
#include "c_QuickFindUF.h"

#define TC_SENTINEL ((c_size_t)-1)

c_err_t c_BoruvkaMST_Init(c_BoruvkaMST_t* self, c_EdgeWeightedGraph_t* graph, c_Allocator_t* allocator) {
    if (!self || !graph) return C_ERR_PARAM;

    self->allocator = allocator ? *allocator : c_DefaultAllocator;
    self->weight = 0.0;
    c_VertexIdList_Init(&self->mst_edges, 0, allocator);

    if (graph->V == 0 || graph->E == 0) return C_SUCCESS;

    /* 1. Initialize Union-Find to manage tree fragments */
    c_QuickFindUF_t uf;
    c_err_t err = c_QuickFindUF_Init(&uf, graph->V, allocator);
    if (err != C_SUCCESS) return err;

    /* 2. Allocate a tracking array to store the closest/cheapest edge ID for each component */
    c_size_t* closest_edge_to_component = (c_size_t*)c_Allocator_Calloc(&self->allocator, graph->V, sizeof(c_size_t));
    if (!closest_edge_to_component) {
        c_QuickFindUF_Destroy(&uf);
        return C_ERR_NOMEM;
    }

    /* 3. Execute successive merging phases (At most log(V) passes) */
    /* Loop until either the tree reaches V - 1 edges or we can't find any more merging cuts */
    for (c_size_t stage = 1; stage < graph->V; stage *= 2) {

        /* Reset closest edge array entries before scanning the edge pool */
        for (c_size_t i = 0; i < graph->V; ++i) {
            closest_edge_to_component[i] = TC_SENTINEL;
        }

        /* Scan every undirected edge to discover the absolute minimum weight cuts for all existing components */
        for (c_size_t e = 0; e < graph->E; ++e) {
            c_Edge_t* edge = &graph->edges_pool[e];
            c_size_t v = edge->v;
            c_size_t w = edge->w;

            c_size_t comp_v = 0;
            c_size_t comp_w = 0;
            c_QuickFindUF_Find(&uf, v, &comp_v);
            c_QuickFindUF_Find(&uf, w, &comp_w);

            /* If they already share the same component ID, adding this edge would create a cycle */
            if (comp_v == comp_w) continue;

            /* Check and update the closest edge for component v */
            if (closest_edge_to_component[comp_v] == TC_SENTINEL ||
                edge->weight < graph->edges_pool[closest_edge_to_component[comp_v]].weight) {
                closest_edge_to_component[comp_v] = e;
            }

            /* Check and update the closest edge for component w */
            if (closest_edge_to_component[comp_w] == TC_SENTINEL ||
                edge->weight < graph->edges_pool[closest_edge_to_component[comp_w]].weight) {
                closest_edge_to_component[comp_w] = e;
            }
        }

        /* Collect and unify components using the best edges found in this phase */
        c_bool_t edges_added_this_phase = C_FALSE;
        for (c_size_t i = 0; i < graph->V; ++i) {
            c_size_t e = closest_edge_to_component[i];
            if (e != TC_SENTINEL) {
                c_Edge_t* edge = &graph->edges_pool[e];
                c_size_t v = edge->v;
                c_size_t w = edge->w;

                if (!c_QuickFindUF_IsConnected(&uf, v, w)) {
                    c_QuickFindUF_Union(&uf, v, w);
                    c_VertexIdList_Append(&self->mst_edges, (c_uint_t)e);
                    self->weight += edge->weight;
                    edges_added_this_phase = C_TRUE;
                }
            }
        }

        /* If no new cross-component cuts were discovered, the spanning tree or forest is complete */
        if (!edges_added_this_phase) break;
    }

    c_Allocator_Free(&self->allocator, closest_edge_to_component);
    c_QuickFindUF_Destroy(&uf);
    return C_SUCCESS;
}

void c_BoruvkaMST_Destroy(c_BoruvkaMST_t* self) {
    if (!self) return;
    c_VertexIdList_Destroy(&self->mst_edges);
    self->weight = 0.0;
}

c_err_t c_BoruvkaMST_GetEdges(c_BoruvkaMST_t* self, c_VertexIdList_t* out_edges) {
    if (!self || !out_edges) return C_ERR_PARAM;
    return c_VertexIdList_Copy(out_edges, &self->mst_edges);
}