//
// Created by Tiago Batista Cardoso on 2/26/2026.
//

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "algorithms.h"

static int has_edge(const graph_t *graph, int u, int v)
{
	node_t *n = graph->adj_lists[u];
	while (n) {
		if (n->id == v)
			return 1;
		n = n->next;
	}
	return 0;
}

typedef struct {
	int **list;
	int count;
	int capacity;
	int k;
} clique_store_t;

static void store_clique(clique_store_t *store, int *current)
{
	if (store->count >= store->capacity) {
		store->capacity *= 2;
		store->list =
			realloc(store->list, store->capacity * sizeof(int *));
	}
	int *copy = malloc(store->k * sizeof(int));
	memcpy(copy, current, store->k * sizeof(int));
	store->list[store->count++] = copy;
}

static void enumerate_cliques(const graph_t *graph, clique_store_t *store,
			      int *current, int depth, int start)
{
	if (depth == store->k) {
		store_clique(store, current);
		return;
	}
	for (int v = start; v < graph->n; v++) {
		// v must be connected to all nodes already in current
		int ok = 1;
		for (int i = 0; i < depth; i++) {
			if (!has_edge(graph, current[i], v)) {
				ok = 0;
				break;
			}
		}
		if (!ok)
			continue;
		current[depth] = v;
		enumerate_cliques(graph, store, current, depth + 1, v + 1);
	}
}

static int uf_find(int *parent, int x)
{
	while (parent[x] != x) {
		parent[x] = parent[parent[x]];
		x = parent[x];
	}
	return x;
}

static void uf_union(int *parent, int *rank, int a, int b)
{
	a = uf_find(parent, a);
	b = uf_find(parent, b);
	if (a == b)
		return;
	if (rank[a] < rank[b]) {
		int t = a;
		a = b;
		b = t;
	}
	parent[b] = a;
	if (rank[a] == rank[b])
		rank[a]++;
}

community_result_t *find_k_clique_communities(const graph_t *graph, int k)
{
	int n = graph->n;

	// find all k-cliques
	clique_store_t store = { .list = malloc(64 * sizeof(int *)),
				 .count = 0,
				 .capacity = 64,
				 .k = k };
	int *current = malloc(k * sizeof(int));
	enumerate_cliques(graph, &store, current, 0, 0);
	free(current);

	printf("[communities] found %d %d-cliques\n", store.count, k);

	int *parent = malloc(store.count * sizeof(int));
	int *rank = calloc(store.count, sizeof(int));
	for (int i = 0; i < store.count; i++)
		parent[i] = i;

	for (int i = 0; i < store.count; i++) {
		for (int j = i + 1; j < store.count; j++) {
			// Count shared nodes
			int shared = 0;
			for (int a = 0; a < k; a++)
				for (int b = 0; b < k; b++)
					if (store.list[i][a] ==
					    store.list[j][b])
						shared++;
			if (shared >= k - 1)
				uf_union(parent, rank, i, j);
		}
	}

	community_result_t *result = malloc(sizeof(community_result_t));
	result->node_community = malloc(n * sizeof(int));
	for (int i = 0; i < n; i++)
		result->node_community[i] = -1;

	for (int i = 0; i < store.count; i++) {
		int community_id = uf_find(parent, i);
		for (int j = 0; j < k; j++) {
			int node = store.list[i][j];
			result->node_community[node] = community_id;
		}
	}

	int *id_map = malloc(store.count * sizeof(int));
	memset(id_map, -1, store.count * sizeof(int));
	int next_id = 0;
	for (int i = 0; i < n; i++) {
		int c = result->node_community[i];
		if (c == -1)
			continue;
		if (id_map[c] == -1)
			id_map[c] = next_id++;
		result->node_community[i] = id_map[c];
	}
	result->count = next_id;
	printf("[communities] found %d communities\n", result->count);

	// cleanup
	for (int i = 0; i < store.count; i++)
		free(store.list[i]);
	free(store.list);
	free(parent);
	free(rank);
	free(id_map);

	return result;
}

void free_community_result(community_result_t *result)
{
	if (!result)
		return;
	free(result->node_community);
	free(result);
}