#include "lib/vec.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

typedef struct HuffCode {
    char byte;
    char *code;
    int frequency;
} HuffCode;

bool huffcode_cmp(HuffCode *first, HuffCode *second) {
    int lendiff = strlen(first->code) - strlen(second->code);
    if (lendiff == 0) {
        int first_code = strtol(first->code, NULL, 2);
        int second_code = strtol(second->code, NULL, 2);
        return first_code >= second_code;
    } else if (lendiff > 0) {
        return true;
    } else {
        return false;
    }
}

DECL_VEC(HuffVec, HuffCode *);

typedef struct MinHeap {
    HuffVec *arr;
} MinHeap;

int minheap_parent(int i) { return (i - 1) / 2; }

int minheap_left_child(int i) { return (2 * i + 1); }

int minheap_right_child(int i) { return (2 * i + 2); }

HuffCode *minheap_get_min(MinHeap *heap) { return heap->arr->data[0]; }

MinHeap *minheap_new(HuffVec *arr) {
    MinHeap *minheap = (MinHeap *)malloc(sizeof(MinHeap));
    minheap->arr = arr;
    return minheap;
}

void minheap_swap(MinHeap *heap, int idx, int jdx) {
    HuffCode *temp = heap->arr->data[idx];
    heap->arr->data[idx] = heap->arr->data[jdx];
    heap->arr->data[jdx] = temp;
}

void *minheap_insert(MinHeap **_heap, HuffCode *value) {
    MinHeap *heap = (*_heap);
    int curr = heap->arr->elem_count - 1;
    HuffVec_push(heap->arr, value);
    while (curr > 0 && huffcode_cmp(heap->arr->data[minheap_parent(curr)],
                                    heap->arr->data[curr])) {
        minheap_swap(heap, minheap_parent(curr), curr);
        curr = minheap_parent(curr);
    }
    return heap;
}

void minheap_print(MinHeap *heap) {
    for (int i = 0; i < heap->arr->elem_count; i++) {
        HuffCode *huffcode = heap->arr->data[i];
        printf("%c: %s\n", huffcode->byte, huffcode->code);
    }
}

int main() {
    printf("Price Hiller, zfp106 (Spring 2024)\n");
    printf("==================================\n");

    HuffVec arr = HuffVec_new(2);
    MinHeap *heap = minheap_new(&arr);

    char chars[] = {'a', 'b', 'c', 'd', 'e'};
    char *codes[] = {"100", "101", "00", "01", "11"};
    int frequencies[] = {5, 9, 12, 13, 16};

    printf("Huffman Encoding (Variable Bit)\n");
    printf("Char | Freq\n");

    HuffCode huffs[sizeof(frequencies) / sizeof(frequencies[0])] = {};
    // Two for loops like this to maintain ownership over the huffman codes
    for (int i = 0; i < sizeof(frequencies) / sizeof(frequencies[0]); i++) {
        HuffCode huff = {
            .byte = chars[i], .code = codes[i], .frequency = frequencies[i]};
        huffs[i] = huff;
        printf("%c    | %d\n", huff.byte, huff.frequency);
    }
    printf("----------\n");

    for (int i = 0; i < sizeof(frequencies) / sizeof(frequencies[0]); i++) {
        heap = minheap_insert(&heap, &huffs[i]);
    }

    minheap_print(heap);
    return 0;
}