#include "aoc.h"
#include <stdlib.h>

typedef struct Node {
    str label;
    i32 index;
} Node;

// Represent edges of graph by an adjacency matrix.
typedef struct BitMatrix {
    i32 size;
    u64 *data;
} BitMatrix;

typedef DYNAMIC_ARRAY(Node) Nodes;

typedef struct Graph {
    Nodes nodes;
    BitMatrix edges;
} Graph;

typedef struct BitSet {
    i32 size;
    i32 array_size;
    u64 *data;
} BitSet;

typedef struct BitSetIter {
    BitSet *bs;
    i32 u64_index;
    u64 current_u64;
} BitSetIter;

static void
bm_init(Arena *arena, BitMatrix *bm, i32 size) {
    bm->size = size;
    isize alloc_size = (bm->size * bm->size + 63) / 64;
    bm->data = ARENA_ALLOC_ARRAY(arena, u64, alloc_size);
}

static u8
bm_at(BitMatrix *bm, i32 row, i32 col) {
    i32 u64_index = (row * bm->size + col) / 64;
    i32 bit_index = (row * bm->size + col) % 64;
    u64 mask = 1ULL << bit_index;
    return (bm->data[u64_index] & mask) != 0;
}

static void
bm_set(BitMatrix *bm, i32 row, i32 col, u8 bit) {
    i32 u64_index = (row * bm->size + col) / 64;
    i32 bit_index = (row * bm->size + col) % 64;
    u64 mask = 1ULL << bit_index;
    if (bit) {
        bm->data[u64_index] |= mask;
    }
    else {
        bm->data[u64_index] &= ~mask;
    }
}

BitSet make_bitset(Arena *arena, i32 size) {
    isize array_size = (size + 63) / 64;
    BitSet set = {
        .size = size,
        .array_size = array_size,
        .data = ARENA_ALLOC_ARRAY(arena, u64, array_size)
    };
    return set;
}

static bool
bs_empty(BitSet *bs) {
    for (i32 i = 0; i < (bs->size + 63) / 64; i++) {
        if (bs->data[i] != 0) {
            return false;
        }
    }
    return true;
}

static void
bs_set(BitSet *bs, i32 index) {
    i32 u64_index = index / 64;
    i32 bit_index = index % 64;
    bs->data[u64_index] |= 1ULL << bit_index;
}

static void
bs_clear(BitSet *bs, i32 index) {
    i32 u64_index = index / 64;
    i32 bit_index = index % 64;
    bs->data[u64_index] &= ~(1ULL << bit_index);
}

static u32
bs_get(BitSet *bs, i32 index) {
    i32 u64_index = index / 64;
    i32 bit_index = index % 64;
    return (bs->data[u64_index] & (1ULL << bit_index)) != 0;
}

BitSetIter
bs_iter(BitSet *bs) {
    BitSetIter iter = {
        .bs = bs,
        .u64_index = -1,
        .current_u64 = 0,
    };
    return iter;
}

static bool
bsi_next(BitSetIter *iter, i32 *index) {
    while (iter->current_u64 == 0) {
        iter->u64_index++;
        if (iter->u64_index >= iter->bs->array_size) {
            return false;
        }
        iter->current_u64 = iter->bs->data[iter->u64_index];
    }
    i32 bit_index = __builtin_ctzll(iter->current_u64);
    *index = iter->u64_index * 64 + bit_index;
    iter->current_u64 &= ~(1ULL << bit_index);
    return true;
}

static BitSet
bs_copy(Arena *arena, BitSet *bs) {
    BitSet copy = make_bitset(arena, bs->size);
    memcpy(copy.data, bs->data, bs->array_size * sizeof(bs->data[0]));
    return copy;
}

static BitSet
bs_intersect(Arena *arena, BitSet *a, BitSet *b) {
    BitSet result = make_bitset(arena, a->size);
    for (i32 i = 0; i < a->array_size; i++) {
        result.data[i] = a->data[i] & b->data[i];
    }
    return result;
}

static i32
bs_count(BitSet *bs) {
    i32 count = 0;
    for (i32 i = 0; i < bs->array_size; i++) {
        count += __builtin_popcountll(bs->data[i]);
    }
    return count;
}

static i32
cmp_i32(const void *p1, const void *p2) {
    i32 a = *(i32 *)p1;
    i32 b = *(i32 *)p2;
    return a - b;
}

static inline bool
graph_connected(Graph *graph, i32 a, i32 b) {
    bool result = bm_at(&graph->edges, a, b);
    return result;
}

static Node *
graph_get_or_insert(Arena *arena, Graph *graph, str label) {
    Node *result = NULL;
    for (isize i = 0; i < graph->nodes.len; i++) {
        if (str_eq(graph->nodes.data[i].label, label)) {
            result = &graph->nodes.data[i];
            break;
        }
    }
    if (!result) {
        Node node = { .label = label, .index = graph->nodes.len };
        *push(&graph->nodes, arena) = node;
        result = &graph->nodes.data[graph->nodes.len - 1];
    }
    return result;
}

static i32
str_cmp(str a, str b) {
    i32 res = memcmp(a.data, b.data, MIN(a.len, b.len));
    if (res == 0) {
        return a.len - b.len;
    }
    else {
        return res;
    }
}

static i32
cmp_str(const void *p1, const void *p2) {
    str *s1 = (str *)p1;
    str *s2 = (str *)p2;
    return str_cmp(*s1, *s2);
}

static BitSet
get_neighbors(Arena *arena, Graph *graph, i32 node) {
    BitSet neighbors = make_bitset(arena, graph->nodes.len);
    for (i32 i = 0; i < graph->nodes.len; i++) {
        if (graph_connected(graph, node, i)) {
            bs_set(&neighbors, i);
        }
    }
    return neighbors;
}

typedef void(CliqueCallback)(BitSet *set, bool maximal, void *context);

static void
bron_kerbosch(Arena *arena, Graph *graph, BitSet R, BitSet P, BitSet X, bool only_maximal, CliqueCallback *callback, void *context) {
    if (bs_empty(&P) && bs_empty(&X)) {
        if (callback) {
            callback(&R, false, context);
        }
    }
    else {
        if (!only_maximal) {
            if (callback) {
                callback(&R, false, context);
            }
        }
        BitSet P_copy = bs_copy(arena, &P);
        BitSetIter iter = bs_iter(&P_copy);
        i32 node;

        void *watermark = arena->top;
        while (bsi_next(&iter, &node)) {
            bs_set(&R, node);

            BitSet neighbors = get_neighbors(arena, graph, node);
            BitSet P_intersect_v = bs_intersect(arena, &P, &neighbors);
            BitSet X_intersect_v = bs_intersect(arena, &X, &neighbors);

            bs_intersect(arena, &P, &P_intersect_v);
            bron_kerbosch(arena, graph, R, P_intersect_v, X_intersect_v, only_maximal, callback, context);
            // undo
            bs_clear(&R, node);

            bs_clear(&P, node);
            bs_set(&X, node);
        }
        arena->top = watermark;
    }
}

typedef struct {
    i32 count;
    Graph *graph;
} Part1_Context;

static void
clique3_callback(BitSet *set, bool maximal, void *context) {
    Part1_Context *p1 = context;
    if (bs_count(set) == 3) {
        i32 node = 0;
        BitSetIter iter = bs_iter(set);
        bool starts_with_t = false;
        while (bsi_next(&iter, &node)) {
            Node *n = &p1->graph->nodes.data[node];
            if (n->label.data[0] == 't') {
                starts_with_t = true;
                break;
            }
        }
        p1->count += starts_with_t;
    }
}

static i32
count_cliques_of_size_3_with_t(Arena *arena, Graph *graph) {
    BitSet P = make_bitset(arena, graph->nodes.len);
    BitSet R = make_bitset(arena, graph->nodes.len);
    BitSet X = make_bitset(arena, graph->nodes.len);

    // create set of all nodes which start with a 't'
    for (i32 i = 0; i < graph->nodes.len; i++) {
        bs_set(&P, i);
    }

    Part1_Context context = {
        .count = 0,
        .graph = graph
    };
    bron_kerbosch(arena, graph, R, P, X, false, clique3_callback, &context);
    return context.count;
}

static void
largest_clique_callback(BitSet *set, bool maximal, void *context) {
    BitSet *largest = context;
    if (bs_count(set) > bs_count(largest)) {
        memcpy(largest->data, set->data, set->array_size * sizeof(set->data[0]));
    }
}

static void
find_largest_clique(Arena *arena, Graph *graph) {
    BitSet P = make_bitset(arena, graph->nodes.len);
    BitSet R = make_bitset(arena, graph->nodes.len);
    BitSet X = make_bitset(arena, graph->nodes.len);

    for (i32 i = 0; i < graph->nodes.len; i++) {
        bs_set(&P, i);
    }

    BitSet result = make_bitset(arena, graph->nodes.len);
    bron_kerbosch(arena, graph, R, P, X, true, largest_clique_callback, &result);

    i32 count = bs_count(&result);
    str *labels = ARENA_ALLOC_ARRAY(arena, str, count);
    i32 node;
    BitSetIter iter = bs_iter(&result);
    i32 counter = 0;
    while (bsi_next(&iter, &node)) {
        labels[counter++] = graph->nodes.data[node].label;
    }
    qsort(labels, count, sizeof(labels[0]), cmp_str);
    for (i32 i = 0; i < count; i++) {
        printf(STR_FMT ",", STR_ARG(labels[i]));
    }
}

int main(int argc, char **argv) {
    Arena *arena = make_arena(Megabytes(2));
    Tokens input = read_lines(arena, argv[1]);
    /* printf("# of edges: %td\n", input.len); */

    Graph graph = {0};
    bm_init(arena, &graph.edges, input.len);

    for (isize i = 0; i < input.len; i++) {
        str line = str_trim(input.tokens[i]);
        if (line.len == 0) {
            continue;
        }
        str a, b;
        str_split(line, STR("-"), &a, &b);

        Node *node_a = graph_get_or_insert(arena, &graph, a);
        Node *node_b = graph_get_or_insert(arena, &graph, b);
        bm_set(&graph.edges, node_a->index, node_b->index, 1);
        bm_set(&graph.edges, node_b->index, node_a->index, 1);
    }

    // Part 1
    i32 result = count_cliques_of_size_3_with_t(arena, &graph);
    printf("%d\n", result);

    // Part 2
    find_largest_clique(arena, &graph);
}