#include <stdio.h>
#include <stdlib.h>

int typedef BOOL;
#define TRUE            1
#define FALSE           0

#define kNYT            257
#define kEmpty          258
#define kMaxTreeHeight  100
#define kRootNumber     51

struct {
    short root;
    short weight;
    short number;
    void *lsub,*rsub;   //pointers to tree's children
    void *parent;       //using a parent pointer quickens things a bit
} typedef tree;

int initialize_tree(tree* tr, unsigned char init_char);
int get_stream_byte(void);
void put_stream_byte(unsigned char byte);
int get_stream_bit(void);
void put_stream_bit(short bit);
int spawn_children(tree* tr, unsigned char init_char);
tree* find_highest(int general_weight, int this_number);
tree* search_tree(short value);
void output_retrace(tree* tr);

unsigned char mask_table[8] = {0x01,0x03,0x07,0x0F,0x1F,0x3F,0x7F,0xFF};
tree *TR;
int global_offset = 7,global_out_offset = 7;
int current_char,current_out_char = 0;
#define NODE_MAX 4096
tree* node_array[4096];  //keep a ongoing list of tree pointers for quick searching
FILE* inFile,*outFile;
BOOL compress_mode;

int main(int argc, char *argv[]) {
    tree* temptr,*max;
    int temp,i = 0,count = 0,this_char;
    
    if (argc < 2 || argc > 4) {
        puts("=======================================================\nUsage:   huffman <source file> <target file> [mode]\n\n   mode:  -c : compress source file to target file\n          -d : decompress source file to target file\n\nBy Spencer Putt\n=======================================================");
        return 0;
    } else if (argc == 3) {
        puts("**Error**   Incorrect number of arguments.  See Usage. **");
        return 0;
    } else { 
        puts("Spencer Putt's Huffman Compressor/Decompressor\n\tNovember 7th, 2005");
        puts("Opening files ... ");
    }
    
    if (argv[3][1] == 'c' || argv[3][1] == 'C') {
        inFile = fopen(argv[1],"rb");
        if (!inFile) {
            puts("Unable to open source file");
            return 1;
        }
        compress_mode = TRUE;
        puts("Starting compression ...");
    } else {
        inFile = fopen(argv[1],"rb");
        if (!inFile) {
            puts("Unable to open source file");
            return 1;
        }
        compress_mode = FALSE;
        puts("Starting decompression ... ");
    }
    if (compress_mode) {
        outFile = fopen(argv[2],"wb");
    } else {
        outFile = fopen(argv[2],"wb");
    }
    if (!outFile) {
        puts("Unable to open destination file");
        return 1;
    }
    current_char = fgetc(inFile);
    this_char = get_stream_byte();
    memset(node_array,0,4096*sizeof(tree*));
    TR = malloc(sizeof(tree));
    if (!TR) {
        puts("Couldn't allocate memory.");
        return 1;
    }
    if (initialize_tree(TR, this_char)) { //initialize the tree with the first char
        puts("Error initializing tree.");
        return 1;
    }
    put_stream_byte((unsigned char) this_char);
    do {
        temptr = TR;
        if (compress_mode) {
            this_char = get_stream_byte();
            if (this_char != EOF) {
                temptr = search_tree(this_char);
                if (temptr) {
                    output_retrace(temptr);
                    // since we got results from search, it means it found
                    // a node which has the correct code.  Now move to the
                    // updater, which will treat this accordingly.
                } else {
                    // we didn't find it, so output kNYT's code and its code.
                    temptr = search_tree(kNYT);
                    output_retrace(temptr); //post the path to the NYT node
                    // in this case, temptr points to kNYT, so new nodes will
                    // spawn correctly.
                }
            }
        } else {
            while (temptr->rsub && temptr->lsub && this_char != EOF) {
                this_char = get_stream_bit();
                if (this_char != EOF) {
                    if (this_char) {
                        temptr = (tree*) temptr->rsub;
                    } else {
                        temptr = (tree*) temptr->lsub;
                    }
                }
            }
        }
        //if the code is NYT, a new char is coming
        if (temptr->root == kNYT) {
            if (!compress_mode) this_char = get_stream_byte();
            if (this_char != EOF) {
                if (spawn_children(temptr, this_char)) {
                    puts("Error allocating memory for new nodes.");
                    return 1;
                }
                temptr->root = kEmpty;      //change tree from kNYT to kEmpty
                if (compress_mode) {
                    put_stream_byte(this_char);
                } else {
                    fputc((int) ((tree*)temptr->rsub)->root, outFile);
                }
            }
        } else {
            if (!compress_mode) {
                fputc((int) temptr->root, outFile);
            }
        }   
        if (this_char == EOF) break;
        do {
            max = find_highest( temptr->weight, temptr->number);
            if (max && (max != (tree*) temptr->parent)) {
                short tempshort;
                void *tempptr;
                // swap children (if they exist)
                if (max->lsub) ((tree*) max->lsub)->parent = temptr;
                if (max->rsub) ((tree*) max->rsub)->parent = temptr;
                if (temptr->lsub) ((tree*) temptr->lsub)->parent = max;
                if (temptr->rsub) ((tree*) temptr->rsub)->parent = max;
                // swap roots
                tempshort = max->root;
                max->root = temptr->root;
                temptr->root = tempshort;
                // swap pointers
                tempptr = max->lsub;
                max->lsub = temptr->lsub;
                temptr->lsub = tempptr;
    
                tempptr = max->rsub;
                max->rsub = temptr->rsub;
                temptr->rsub = tempptr;
    
                temptr = max;
            }
            temptr->weight++;
            temptr = temptr->parent; 
        } while (temptr != NULL);       //root node will have no parent
    } while (this_char != EOF);

    if (compress_mode) {
        temptr = search_tree(kNYT);
        output_retrace(temptr);    //fill the rest of the byte with useless node
        puts("Compression completed successfully.");
    } else {
        puts("Decompression completed.");
    }
    fclose(inFile);
    fclose(outFile);

    return 0;
}

//initializes base tree based on first byte in stream
int initialize_tree(tree* tr, unsigned char init_char) {
    tr->root = kEmpty;
    tr->weight = 1;
    tr->number = kRootNumber;
    tr->parent = NULL;
    node_array[0] = tr;
    return spawn_children(tr, init_char);
}

//creates a new NYT node and node containing new value
int spawn_children(tree* tr, unsigned char init_char) {
    int i;
    tr->lsub = malloc(sizeof(tree));
    tr->rsub = malloc(sizeof(tree));
    if (!tr->lsub || !tr->rsub) {
        puts("Could not allocate memory");
        return -1;
    }
    for (i = 0; i < 4096, node_array[i] != 0; i++);
    node_array[i++] = tr->lsub;
    node_array[i] = tr->rsub;

    ((tree*) tr->rsub)->root = (short) init_char;
    ((tree*) tr->rsub)->weight = 1;
    ((tree*) tr->rsub)->number = tr->number - 1;
    ((tree*) tr->rsub)->lsub = NULL;
    ((tree*) tr->rsub)->rsub = NULL;
    ((tree*) tr->rsub)->parent = tr;

    ((tree*) tr->lsub)->root = kNYT;
    ((tree*) tr->lsub)->weight = 0;
    ((tree*) tr->lsub)->number = tr->number - 2;
    ((tree*) tr->lsub)->lsub = NULL;
    ((tree*) tr->lsub)->rsub = NULL;
    ((tree*) tr->lsub)->parent = tr;
    return 0;
}

//gets a bit from the input stream
int get_stream_bit(void) {
    unsigned char mask = 0x01;
    int temp;
    if (feof(inFile)) return EOF;
    if (global_offset) mask = mask << global_offset;
    temp = current_char & mask;
    if (global_offset == 0) {
        global_offset = 8;
        current_char = fgetc(inFile);
    }
    global_offset--;
    if (temp) return 1;
    return 0;
}

//puts a bit on the output stream
void put_stream_bit(short bit) {
    unsigned char mask = 0x01;
    if (bit) {
        if (global_out_offset) mask = mask << global_out_offset;
        current_out_char |= mask;
    }
    if (global_out_offset == 0) {
        global_out_offset = 8;
        fputc((int) current_out_char, outFile);
        current_out_char = 0;
    }
    global_out_offset--;
}

//gets a byte from the input stream
int get_stream_byte(void) {
    unsigned char mask;
    int final = 0;
    if (current_char == EOF) return EOF;
    mask = mask_table[global_offset];
    final = (current_char & mask)<<8;
    current_char = fgetc(inFile);
    final += (current_char & ((unsigned char) ~mask));
    final = final >> (global_offset+1);
      
    int loc = ftell(inFile);
    if (loc % 1000 == 0) printf("loc: %d\n",loc);  
    return final;

}

//puts a byte (aligned or not) on the output stream
void put_stream_byte(unsigned char byte) {
    unsigned char mask = byte;
    if (global_out_offset != 7) mask = byte >> (7 - global_out_offset);
    current_out_char |= mask;
    fputc(current_out_char, outFile);
    current_out_char = 0;
    mask = 0;
    if (global_out_offset != 7) mask = byte << (1 + global_out_offset);
    current_out_char |= mask;
}

// returns pointer to highest tree element of given weight, 
// NULL if current element is highest
tree* find_highest(int general_weight, int this_number) {
    tree* max = NULL;
    int i;
    for (i = 0; i < 4096; i ++) {
        if (node_array[i]) {
            if ((node_array[i]->weight == general_weight) 
            && (node_array[i]->number > this_number)) {
                max = node_array[i];
                this_number = max->number;
            }
        }
    }
    return max;
}

//searches tree roots for a given value
tree* search_tree(short value) {
    int i;
    for (i = 0; i < 4096; i++) {
        if (node_array[i]) {
            if (node_array[i]->root == value) return node_array[i];
        }
    }
    return NULL;
}

//traces path back to tree and outputs it to outFile;
void output_retrace(tree* tr) {
    int i,bit_list[4096];
    tree* temp;
    
    for (i = 0; tr != TR; i++) {
        temp = tr->parent;
        if (tr == (tree*) temp->lsub) {
            bit_list[i] = 0;
        } else {
            bit_list[i] = 1;
        }
        tr = temp;
    }
    i--;
    while (i >= 0) {
        put_stream_bit((short) bit_list[i--]);
    }
}