#include "huffman.h"
#include <stdlib.h>
#include <string.h>


/********************************************************
 Audio Tools, a module and set of tools for manipulating audio data
 Copyright (C) 2007-2014  Brian Langenberger

 This program is free software; you can redistribute it and/or modify
 it under the terms of the GNU General Public License as published by
 the Free Software Foundation; either version 2 of the License, or
 (at your option) any later version.

 This program is distributed in the hope that it will be useful,
 but WITHOUT ANY WARRANTY; without even the implied warranty of
 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 GNU General Public License for more details.

 You should have received a copy of the GNU General Public License
 along with this program; if not, write to the Free Software
 Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
*******************************************************/

typedef enum {NODE_TREE, NODE_LEAF} huffman_node_type;

struct huffman_node {
    huffman_node_type type;
    union {
        int leaf;
        struct {
            unsigned int id;
            br_huffman_table_t jump_table;
            struct huffman_node* bit_0;
            struct huffman_node* bit_1;
        } tree;
    } v;
};

#define BYTE_BANK_SIZE 9

struct byte_bank {
    unsigned int size;
    unsigned int value;
};

#define CONTINUE_READING (1 << BYTE_BANK_SIZE)

/*takes a list Huffman frequencies and returns a completed Huffman tree

  the final huffman_frequency in the list should have a length of 0
  and serve as our list terminator

  the tree must be freed when no longer needed*/
static struct huffman_node*
build_huffman_tree(struct huffman_frequency* frequencies,
                   unsigned int total_frequencies,
                   int* error);

/*the proper recursive version, not to be called directly most of the time*/
static struct huffman_node*
build_huffman_tree_(unsigned int bits,
                    unsigned int length,
                    struct huffman_frequency* frequencies,
                    unsigned int total_frequencies,
                    unsigned int* counter,
                    int* error);

/*deallocates the space for the Huffman tree*/
static void
free_huffman_tree(struct huffman_node* node);

/*takes a built Huffman tree and compiles it to a jump table
  with the given endianness*/
static int
compile_huffman_tree(br_huffman_table_t** table,
                     struct huffman_node* tree,
                     bs_endianness endianness);

/*populates the jump tables of the Huffman tree's nodes*/
static void
populate_huffman_tree(struct huffman_node* tree,
                      bs_endianness endianness);

/*returns the total number of non-leaf nodes in the tree*/
static unsigned int
total_non_leaf_nodes(struct huffman_node* tree);

/*returns the total number of leaf nodes in the tree*/
static unsigned int
total_leaf_nodes(struct huffman_node* tree);

/*transfers the jump tables embedded in the Huffman tree nodes
  to a flattened, 2 dimensional array*/
static void
transfer_huffman_tree(br_huffman_table_t* table,
                      struct huffman_node* tree);

/*performs the actual walking of the tree

  given a byte bank, endianness and tree,
  updates the jump table with which node the bank will wind up at*/
static void
next_read_huffman_state(br_huffman_entry_t* state,
                        struct byte_bank bank,
                        struct huffman_node* tree,
                        bs_endianness endianness);

/*converts the byte bank to a regular integer*/
static int
bank_to_int(struct byte_bank bank);


static void
bw_fix_tree(bw_huffman_table_t *table,
            int value_index,
            int value);


/*returns a new set of bits of the same length
  such that the least significant become the most significant
  and vice versa*/
static unsigned int
swap_bits(unsigned int bits, unsigned int length);


static int
bank_to_int(struct byte_bank bank) {
    if (bank.size > 0) {
        assert(bank.value <= ((1 << bank.size) - 1));
        return (1 << bank.size) | bank.value;
    } else
        return 0;
}

static struct huffman_node*
build_huffman_tree(struct huffman_frequency* frequencies,
                   unsigned int total_frequencies,
                   int* error)
{
    unsigned int counter = 0;
    struct huffman_node* built_tree;

    built_tree = build_huffman_tree_(0, 0,
                                     frequencies, total_frequencies,
                                     &counter, error);
    if (built_tree != NULL) {
        if (total_frequencies > total_leaf_nodes(built_tree)) {
            *error = HUFFMAN_ORPHANED_LEAF;
            free_huffman_tree(built_tree);
            return NULL;
        } else {
            return built_tree;
        }
    } else {
        return NULL;
    }
}

static struct huffman_node*
build_huffman_tree_(unsigned int bits,
                    unsigned int length,
                    struct huffman_frequency* frequencies,
                    unsigned int total_frequencies,
                    unsigned int* counter,
                    int* error)
{
    unsigned int i;
    unsigned int j;
    struct huffman_node* node = malloc(sizeof(struct huffman_node));
    unsigned int max_frequency_length = 0;

    /*go through the list of frequency values*/
    for (i = 0; i < total_frequencies; i++) {
        /*if our bits and length value is found,
          generate a new leaf node from that frequency
          so long as it is unique in the list*/
        if ((frequencies[i].bits == bits) &&
            (frequencies[i].length == length)) {
            /*check for duplicates*/
            for (j = i + 1; j < total_frequencies; j++) {
                if ((frequencies[j].bits == bits) &&
                    (frequencies[j].length == length)) {
                    *error = HUFFMAN_DUPLICATE_LEAF;
                    free(node);
                    return NULL;
                }
            }

            node->type = NODE_LEAF;
            node->v.leaf = frequencies[i].value;
            return node;
        } else {
            max_frequency_length = MAX(max_frequency_length,
                                       frequencies[i].length);
        }
    }

    if (length > max_frequency_length) {
        /*we've walked outside of the set of possible frequencies
          which indicates the tree is missing a leaf node*/
        *error = HUFFMAN_MISSING_LEAF;
        free(node);
        return NULL;
    }

    /*otherwise, generate a new tree node
      whose leaf nodes are generated recursively*/
    node->type = NODE_TREE;
    node->v.tree.id = *counter;
    node->v.tree.bit_0 = NULL;
    node->v.tree.bit_1 = NULL;
    (*counter) += 1;
    if ((node->v.tree.bit_0 = build_huffman_tree_(bits << 1,
                                                  length + 1,
                                                  frequencies,
                                                  total_frequencies,
                                                  counter,
                                                  error)) == NULL)
        goto error;

    if ((node->v.tree.bit_1 = build_huffman_tree_((bits << 1) | 1,
                                                  length + 1,
                                                  frequencies,
                                                  total_frequencies,
                                                  counter,
                                                  error)) == NULL)
        goto error;

    return node;
 error:
    free_huffman_tree(node->v.tree.bit_0);
    free_huffman_tree(node->v.tree.bit_1);
    free(node);
    return NULL;
}

static void
free_huffman_tree(struct huffman_node* node) {
    if (node == NULL)
        return;
    else if (node->type == NODE_LEAF) {
        free(node);
    } else {
        free_huffman_tree(node->v.tree.bit_0);
        free_huffman_tree(node->v.tree.bit_1);
        free(node);
    }
}

static int
compile_huffman_tree(br_huffman_table_t** table,
                     struct huffman_node* tree,
                     bs_endianness endianness) {
    int total_rows = total_non_leaf_nodes(tree);
    unsigned int size;
    unsigned int value;
    struct byte_bank bank;
    int bank_int;

    if (total_rows > 0) {
        /*allocate space for all the jump tables across all nodes*/
        br_huffman_table_t *new_table =
            malloc(sizeof(br_huffman_table_t) * total_rows);

        /*populate the jump tables of each non-leaf node*/
        populate_huffman_tree(tree, endianness);

        /*transfer jump tables of each node from tree*/
        transfer_huffman_tree(new_table, tree);

        /*transfer new table back to parent*/
        *table = new_table;
    } else if (total_leaf_nodes(tree) > 0) {
        /*no non-leaf nodes, so the table is trivial
          all inputs consume no bits and return the final value*/

        br_huffman_table_t *new_table =
            malloc(sizeof(br_huffman_table_t) * 1);

        /*first 2 nodes are a special case
          since a context of 0 or 1 has no meaning*/
        new_table[0][0].continue_ = 0;
        new_table[0][0].node = 0;
        new_table[0][0].state = 0;
        new_table[0][0].value = tree->v.leaf;
        new_table[0][1].continue_ = 0;
        new_table[0][1].node = 0;
        new_table[0][1].state = 0;
        new_table[0][1].value = tree->v.leaf;

        for (size = 1; size < (8 + 1); size++)
            for (value = 0; value < (1 << size); value++) {
                bank.size = size;
                bank.value = value;
                bank_int = bank_to_int(bank);

                new_table[0][bank_int].continue_ = 0;
                new_table[0][bank_int].node = 0;
                new_table[0][bank_int].state = bank_int;
                new_table[0][bank_int].value = tree->v.leaf;
            }

        total_rows = 1;
    } else {
        *table = malloc(0);
        return HUFFMAN_EMPTY_TREE;
    }

    return total_rows;
}

static void
populate_huffman_tree(struct huffman_node* tree,
                      bs_endianness endianness) {
    unsigned int size;
    unsigned int value;
    struct byte_bank bank;

    if (tree->type == NODE_TREE) {
        tree->v.tree.jump_table[0].continue_ = 1;
        tree->v.tree.jump_table[0].node = 0;
        tree->v.tree.jump_table[0].state = 0;
        tree->v.tree.jump_table[0].value = 0;
        tree->v.tree.jump_table[1].continue_ = 1;
        tree->v.tree.jump_table[1].node = 0;
        tree->v.tree.jump_table[1].state = 0;
        tree->v.tree.jump_table[1].value = 0;

        for (size = 1; size < (8 + 1); size++)
            for (value = 0; value < (1 << size); value++) {
                bank.size = size;
                bank.value = value;

                next_read_huffman_state(
                    &(tree->v.tree.jump_table[bank_to_int(bank)]),
                    bank, tree, endianness);
        }

        populate_huffman_tree(tree->v.tree.bit_0, endianness);
        populate_huffman_tree(tree->v.tree.bit_1, endianness);
    }
}

void next_read_huffman_state(br_huffman_entry_t* state,
                             struct byte_bank bank,
                             struct huffman_node* tree,
                             bs_endianness endianness) {
    struct byte_bank next_bank;

    if (tree->type == NODE_LEAF) {
        /*reached a leaf node,
          so return current byte bank, empty continue bit and value*/
        state->continue_ = 0;
        state->node = 0;
        state->state = bank_to_int(bank);
        state->value = tree->v.leaf;
    } else if (bank.size == 0) {
        /*exhausted byte bank,
          so return empty bank, set continue bit and current node*/
        state->continue_ = 1;
        state->node = tree->v.tree.id;
        state->state = 0;
        state->value = 0;
    } else if (endianness == BS_LITTLE_ENDIAN) {
        /*progress through bit stream in little endian order*/
        next_bank = bank;
        next_bank.size -= 1;
        next_bank.value >>= 1;

        if (bank.value & 1) {
            next_read_huffman_state(state,
                                    next_bank,
                                    tree->v.tree.bit_1,
                                    endianness);
        } else {
            next_read_huffman_state(state,
                                    next_bank,
                                    tree->v.tree.bit_0,
                                    endianness);

        }
    } else if (endianness == BS_BIG_ENDIAN) {
        /*progress through bit stream in big endian order*/
        next_bank = bank;
        next_bank.size -= 1;
        next_bank.value &= ((1 << next_bank.size) - 1);

        if (bank.value & (1 << (bank.size - 1))) {
            next_read_huffman_state(state,
                                    next_bank,
                                    tree->v.tree.bit_1,
                                    endianness);
        } else {
            next_read_huffman_state(state,
                                    next_bank,
                                    tree->v.tree.bit_0,
                                    endianness);

        }
    }
}

static unsigned int
total_non_leaf_nodes(struct huffman_node* tree) {
    if (tree->type == NODE_TREE) {
        return (1 +
                total_non_leaf_nodes(tree->v.tree.bit_0) +
                total_non_leaf_nodes(tree->v.tree.bit_1));
    } else
        return 0;
}

static unsigned int
total_leaf_nodes(struct huffman_node* tree) {
    if (tree->type == NODE_TREE) {
        return (total_leaf_nodes(tree->v.tree.bit_0) +
                total_leaf_nodes(tree->v.tree.bit_1));
    } else
        return 1;
}

static void
transfer_huffman_tree(br_huffman_table_t* table,
                      struct huffman_node* tree) {
    int i;

    if (tree->type == NODE_TREE) {
        /*not sure if this can be made more efficient*/
        for (i = 0; i < 0x200; i++) {
            table[tree->v.tree.id][i] = tree->v.tree.jump_table[i];
        }
        transfer_huffman_tree(table, tree->v.tree.bit_0);
        transfer_huffman_tree(table, tree->v.tree.bit_1);
    }
}

int compile_br_huffman_table(br_huffman_table_t** table,
                             struct huffman_frequency* frequencies,
                             unsigned int total_frequencies,
                             bs_endianness endianness) {
    int error = 0;
    struct huffman_node* tree;
    int total_rows;

    tree = build_huffman_tree(frequencies, total_frequencies, &error);
    if (tree == NULL)
        return error;
    total_rows = compile_huffman_tree(table, tree, endianness);
    free_huffman_tree(tree);
    return total_rows;
}

typedef int (*qsort_cmp_func_t)(const void*, const void*);

static
int frequency_cmp(const struct huffman_frequency* f1,
                  const struct huffman_frequency* f2)
{
    if (f1->length == f2->length) {
        return 0;
    } else if (f1->length < f2->length) {
        return -1;
    } else {
        return 1;
    }
}

int compile_bw_huffman_table(bw_huffman_table_t** table,
                             struct huffman_frequency* frequencies,
                             unsigned int total_frequencies,
                             bs_endianness endianness)
{
    bw_huffman_table_t *new_table;
    int error = 0;
    struct huffman_node* tree;
    unsigned int i;

    *table = NULL;

    /*ensure frequencies generate a valid Huffman tree*/
    tree = build_huffman_tree(frequencies, total_frequencies, &error);
    if (tree == NULL) {
        return error;
    } else {
        free_huffman_tree(tree);
    }

    /*sort frequencies in ascending order of bit count
      to ensure the most common items are near the top of the tree*/
    qsort(frequencies,
          (size_t)total_frequencies,
          sizeof(struct huffman_frequency),
          (qsort_cmp_func_t)frequency_cmp);

    new_table = malloc(sizeof(bw_huffman_table_t) * total_frequencies);

    /*for each frequency in the list*/
    for (i = 0; i < total_frequencies; i++) {
        /*insert a node into the binary tree
          with the given value, bit count and bit value*/
        new_table[i].value = frequencies[i].value;
        new_table[i].write_count = frequencies[i].length;
        if (endianness == BS_BIG_ENDIAN) {
            new_table[i].write_value = frequencies[i].bits;
        } else {
            new_table[i].write_value = swap_bits(frequencies[i].bits,
                                                 frequencies[i].length);
        }
        new_table[i].smaller = -1;
        new_table[i].larger = -1;

        bw_fix_tree(new_table, i, frequencies[i].value);
    }

    *table = new_table;

    return 0;
}

static unsigned int
swap_bits(unsigned int bits, unsigned int length){
    unsigned int swapped = 0;

    while (length--) {
        const unsigned int lsb = bits & 1;
        bits >>= 1;
        swapped = (swapped << 1) | lsb;
    }

    return swapped;
}

static void
bw_fix_tree(bw_huffman_table_t *table,
            int value_index,
            int value)
{
    int current_index = 0;
    while (current_index != value_index) {
        const int table_value = table[current_index].value;
        if (value == table_value) {
            /*entry with same value found in tree
              so don't bother adding new one*/
            return;
        } else if (value < table_value) {
            if (table[current_index].smaller == -1) {
                table[current_index].smaller = value_index;
            }
            current_index = table[current_index].smaller;
        } else {
            if (table[current_index].larger == -1) {
                table[current_index].larger = value_index;
            }
            current_index = table[current_index].larger;
        }
    }
}

#ifdef EXECUTABLE

#include "parson.h"
#include <getopt.h>

struct huffman_frequency* json_to_frequencies(const char* path,
                                              unsigned int* total_frequencies);

struct huffman_frequency parse_json_pair(JSON_Array* bit_list, double value);

int main(int argc, char* argv[]) {
    /*option handling variables*/
    static int little_endian_arg = 0;
    static int write_arg = 0;
    bs_endianness endianness;
    char* input_file = NULL;

    static struct option long_options[] = {
        {"input", required_argument, 0, 'i'},
        {"help", no_argument, 0, 'h'},
        {"le", no_argument, &little_endian_arg, 1},
        {"write", no_argument, &write_arg, 1},
        {0, 0, 0, 0}
    };

    int option_index = 0;
    int c;

    /*the variables for real work*/
    struct huffman_frequency* frequencies;
    unsigned int total_frequencies;

    do {
        c = getopt_long(argc, argv, "i:h", long_options, &option_index);
        switch (c) {
        case 'h':
            printf("Options:\n");
            printf("  -h, --help             "
                   "show this help message and exit\n");
            printf("  -i PATH, --input=PATH  "
                   "input JSON file\n");
            printf("  --le                   "
                   "generate little-endian jump table\n");
            printf("  --write                "
                   "generate a write table\n");
            return 0;
        case 'i':
            input_file = optarg;
            break;
        case '?':
            return 1;
        case 0:
        case -1:
        default:
            break;
        }
    } while (c != -1);

    if (input_file == NULL) {
        fprintf(stderr, "an input file is required\n");
        return 1;
    }

    endianness = little_endian_arg ? BS_LITTLE_ENDIAN : BS_BIG_ENDIAN;

    frequencies = json_to_frequencies(input_file, &total_frequencies);

    if (!write_arg) {
        br_huffman_table_t* table;
        const int total_rows = compile_br_huffman_table(&table,
                                                        frequencies,
                                                        total_frequencies,
                                                        endianness);
        int row;

        if (total_rows < 0) {
            switch (total_rows) {
            case HUFFMAN_MISSING_LEAF:
                fprintf(stderr, "Huffman table missing leaf node\n");
                free(frequencies);
                return 1;
            case HUFFMAN_DUPLICATE_LEAF:
                fprintf(stderr, "Huffman table has duplicate leaf node\n");
                free(frequencies);
                return 1;
            case HUFFMAN_ORPHANED_LEAF:
                fprintf(stderr, "Huffman table has orphaned leaf nodes\n");
                free(frequencies);
                return 1;
            default:
                fprintf(stderr, "Unknown error\n");
                free(frequencies);
                return 1;
            }
        }

        printf("{\n");
        for (row = 0; row < total_rows; row++) {
            int state;

            printf("  {\n");

            for (state = 0; state < 0x200; state++) {
                printf("    {%d, %d, 0x%X, %d}",
                       table[row][state].continue_,
                       table[row][state].node,
                       table[row][state].state,
                       table[row][state].value);
                if (state < (0x200 - 1))
                    printf(",\n");
                else
                    printf("\n");
            }
            if (row < (total_rows - 1))
                printf("  },\n");
            else
                printf("  }\n");
        }
        printf("}\n");

        free(table);
    } else {
        bw_huffman_table_t *table;

        const int result = compile_bw_huffman_table(&table,
                                                    frequencies,
                                                    total_frequencies,
                                                    endianness);
        int row;

        if (result < 0) {
            switch (result) {
            case HUFFMAN_MISSING_LEAF:
                fprintf(stderr, "Huffman table missing leaf node\n");
                free(frequencies);
                return 1;
            case HUFFMAN_DUPLICATE_LEAF:
                fprintf(stderr, "Huffman table has duplicate leaf node\n");
                free(frequencies);
                return 1;
            case HUFFMAN_ORPHANED_LEAF:
                fprintf(stderr, "Huffman table has orphaned leaf nodes\n");
                free(frequencies);
                return 1;
            default:
                fprintf(stderr, "Unknown error\n");
                free(frequencies);
                return 1;
            }
        }

        printf("{\n");

        for (row = 0; row < total_frequencies; row++) {
            printf("  {%d, %u, 0x%X, %d, %d",
                   table[row].value,
                   table[row].write_count,
                   table[row].write_value,
                   table[row].smaller,
                   table[row].larger);
            if (row < (total_frequencies - 1)) {
                printf("},\n");
            } else {
                printf("}\n");
            }
        }

        printf("}\n");

        free(table);
    }

    free(frequencies);
    return 0;
}

struct huffman_frequency* json_to_frequencies(const char* path,
                                              unsigned int* total_frequencies) {
    JSON_Value* file;
    JSON_Array* input;
    size_t input_size;
    int o;
    size_t i;
    struct huffman_frequency* frequencies;

    if ((file = json_parse_file(path)) == NULL) {
        fprintf(stderr, "error parsing input .json file \"%s\"\n", path);
        exit(1);
    } else {
        input = json_value_get_array(file);
        if (input == NULL) {
            fprintf(stderr, "JSON file isn't an array of items\n");
            exit(1);
        }
    }

    input_size = json_array_get_count(input);

    frequencies = malloc(sizeof(struct huffman_frequency) * (input_size / 2));

    *total_frequencies = input_size / 2;
    for (i = o = 0; i < input_size; i += 2,o++) {
        frequencies[o] = parse_json_pair(json_array_get_array(input, i),
                                         json_array_get_number(input, i + 1));
    }

    json_value_free(file);

    return frequencies;
}

struct huffman_frequency parse_json_pair(JSON_Array* bit_list, double value) {
    struct huffman_frequency frequency;
    size_t i;

    frequency.bits = 0;
    frequency.length = 0;

    for (i = 0; i < json_array_get_count(bit_list); i++) {
        frequency.bits = ((frequency.bits << 1) |
                          (int)(json_array_get_number(bit_list, i)));
        frequency.length++;
    }

    frequency.value = (int)(value);

    return frequency;
}

#endif
