/* c1.c - Don Yang (uguu.org)

   Removed strict types.

   10/11/01
                                                                           */

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

#define INTERNAL_NODE   (-1)
#define UNUSED_NODE     300

typedef struct _node
{
   int freq, key, id, bit;
   struct _node *left, *right, *next;
} Node;

static Node HeapPool[512], *Heap;
static int HeapPoolIndex;

static Node *Dictionary[256];

static int Left[256], Right[256], NodeCount;
static int FileSize;

static void AssignID(Node *node, Node *parent);
static void CreateTree(FILE *infile);
static int DecodeFile(FILE *infile, FILE *outfile);
static int EncodeFile(FILE *infile, FILE *outfile);
static Node *HeapPop(void);
static void HeapPush(int freq, int key, Node *left, Node *right);
static int LoadHeader(unsigned char *buffer, int size);
static int LoadTree(unsigned char *buffer);
static int WriteBits(FILE *outfile, int *byte, int *bit, Node *node);
static int WriteHeader(FILE *outfile);

/******************************************************************** main */
int main(int argc, char **argv)
{
   static unsigned char buffer[579];
   FILE *infile, *outfile;
   size_t size;

   if( argc < 3 )
      return printf("%s <infile> <outfile>\n", *argv);

   if( (infile = fopen(argv[1], "rb")) == NULL )
      return printf("Can not open %s\n", argv[1]);
   if( (outfile = fopen(argv[2], "wb+")) == NULL )
   {
      fclose(infile);
      return printf("Can not create %s\n", argv[2]);
   }

   fseek(infile, 0, SEEK_SET);
   size = fread(buffer, 1, 579, infile);
   if( size == 0 )
   {
      fclose(infile);
      fclose(outfile);
      return puts("Empty file");
   }
   if( LoadHeader(buffer, size) == 0 )
   {
      if( DecodeFile(infile, outfile) != 0 )
         printf("Error writing %s\n", argv[2]);
      else
         printf("Decoded %s -> %s\n", argv[1], argv[2]);
   }
   else
   {
      CreateTree(infile);
      if( WriteHeader(outfile) != 0 )
         printf("Error writing %s\n", argv[2]);
      else if( EncodeFile(infile, outfile) != 0 )
         printf("Error writing %s\n", argv[2]);
      else
         printf("Encoded %s -> %s\n", argv[1], argv[2]);
   }

   fclose(infile);
   fclose(outfile);
   return 0;

} /* main() */

/**************************************************************** AssignID */
static void AssignID(Node *node, Node *parent)
{
   int id;

   node->next = parent;
   if( node->key != INTERNAL_NODE )
   {
      Dictionary[node->key] = node;
      return;
   }
   node->id = id = NodeCount++;
   AssignID(node->left, node);   node->left->bit = 0;
   AssignID(node->right, node);  node->right->bit = 1;

   Left[id] = (node->left->key == INTERNAL_NODE)
              ? node->left->id : -node->left->key;
   Right[id] = (node->right->key == INTERNAL_NODE)
               ? node->right->id : -node->right->key;

} /* AssignID() */

/************************************************************** CreateTree */
static void CreateTree(FILE *infile)
{
   static unsigned char buffer[1024];
   int size, i, freq[256];
   Node *a, *b;

   for(i = 0; i < 256; freq[i++] = 0);

   FileSize = 0;
   fseek(infile, 0, SEEK_SET);
   while( (size = fread(buffer, 1, 1024, infile)) != 0 )
   {
      for(i = 0; i < size; freq[buffer[i++]]++);
      FileSize += size;
   }

   HeapPoolIndex = 0;
   Heap = NULL;

   for(i = 0; i < 256; i++)
   {
      if( freq[i] != 0 )
         HeapPush(freq[i], i, NULL, NULL);
   }
   if( Heap->next == NULL )
   {
      for(i = 0; freq[i] > 0; i++);
      HeapPush(freq[i], i, NULL, NULL);
   }

   do
   {
      a = HeapPop();
      b = HeapPop();
      HeapPush(a->freq + b->freq, INTERNAL_NODE, a, b);
   } while( Heap != NULL && Heap->next != NULL );

   NodeCount = 0;
   AssignID(Heap, NULL);

   for(i = NodeCount; i < 256; i++)
      Left[i] = Right[i] = UNUSED_NODE;

} /* CreateTree() */

/************************************************************** DecodeFile */
static int DecodeFile(FILE *infile, FILE *outfile)
{
   int bit, byte, node, *branch, i;

   fseek(infile, NodeCount * 2 + (NodeCount + 3) / 4 + 5, SEEK_SET);
   byte = node = 0;
   bit = 8;
   for(i = 0; i < FileSize;)
   {
      if( bit > 7 )
      {
         if( (byte = fgetc(infile)) == EOF )
            return 1;
         bit = 0;
      }

      branch = (byte & (1 << bit)) == 0 ? Left : Right;
      if( branch[node] <= 0 )
      {
         if( fputc(-branch[node], outfile) == EOF )
            return 1;
         node = 0;
         i++;
      }
      else
      {
         if( branch[node] <= node )
            return 1;
         node = branch[node];
      }
      bit++;
   }

   return 0;

} /* DecodeFile() */

/************************************************************** EncodeFile */
static int EncodeFile(FILE *infile, FILE *outfile)
{
   int bit, byte, data, i;

   bit = byte = 0;
   fseek(infile, 0, SEEK_SET);
   for(i = 0; i < FileSize; i++)
   {
      if( (data = fgetc(infile)) == EOF )
         return 1;
      if( WriteBits(outfile, &byte, &bit, Dictionary[data]) != 0 )
         return 1;
   }

   if( bit > 0 )
   {
      if( fputc(byte, outfile) == EOF )
         return 1;
   }
   return 0;

} /* EncodeFile() */

/***************************************************************** HeapPop */
static Node *HeapPop(void)
{
   Node *r;

   r = Heap;
   Heap = Heap->next;
   return r;

} /* HeapPop() */

/**************************************************************** HeapPush */
static void HeapPush(int freq, int key, Node *left, Node *right)
{
   Node *node, **cursor;

   node = &HeapPool[HeapPoolIndex++];

   node->freq = freq;
   node->key = key;
   node->left = left;
   node->right = right;

   for(cursor = &Heap;
       *cursor != NULL && (*cursor)->freq < node->freq;
       cursor = &((*cursor)->next));
   node->next = *cursor;
   *cursor = node;

} /* HeapPush() */

/************************************************************** LoadHeader */
static int LoadHeader(unsigned char *buffer, int size)
{
   NodeCount = buffer[0];
   if( NodeCount * 2 + ((NodeCount + 3) / 4) + 5 > size )
      return 1;
   return LoadTree(buffer);

} /* LoadHeader() */

/**************************************************************** LoadTree */
static int LoadTree(unsigned char *buffer)
{
   int leaf[256], i, j;

   for(i = 0; i < NodeCount; i += 4)
   {
      for(j = 0; j < 4; j++)
         leaf[i + j] = ((buffer[i / 4 + 1] >> (j * 2)) & 3);
   }

   j = (NodeCount + 3) / 4 + 1;
   for(i = 0; i < NodeCount; i++)
   {
      if( (leaf[i] & 1) != 0 )
      {
         Left[i] = -buffer[j++];
      }
      else
      {
         Left[i] = buffer[j++];
         if( Left[i] >= NodeCount || Left[i] <= i )
            return 1;
      }
      if( (leaf[i] & 2) != 0 )
      {
         Right[i] = -buffer[j++];
      }
      else
      {
         Right[i] = buffer[j++];
         if( Right[i] >= NodeCount || Right[i] <= i )
            return 1;
      }
   }
   FileSize = (buffer[j] | (buffer[j + 1] << 8) |
              (buffer[j + 2] << 16) | (buffer[j + 3] << 24));
   return 0;

} /* LoadTree() */

/*************************************************************** WriteBits */
static int WriteBits(FILE *outfile, int *byte, int *bit, Node *node)
{
   if( node->next != NULL )
   {
      if( WriteBits(outfile, byte, bit, node->next) != 0 )
         return 1;
      *byte |= node->bit << (*bit)++;
      if( *bit > 7 )
      {
         if( fputc(*byte, outfile) == EOF )
            return 1;
         *bit = *byte = 0;
      }
   }
   return 0;

} /* WriteBits() */

/************************************************************* WriteHeader */
static int WriteHeader(FILE *outfile)
{
   int i, j, byte;

   if( fputc(NodeCount, outfile) == EOF )
      return 1;
   for(i = 0; i < NodeCount; i += 4)
   {
      byte = 0;
      for(j = 0; j < 4; j++)
      {
         if( Left[i + j] <= 0 )  byte |= 1 << (j * 2);
         if( Right[i + j] <= 0 ) byte |= 2 << (j * 2);
      }
      if( fputc(byte, outfile) == EOF )
         return 1;
   }

   for(i = 0; i < NodeCount; i++)
   {
      if( fputc(abs(Left[i]), outfile) == EOF )    return 1;
      if( fputc(abs(Right[i]), outfile) == EOF )   return 1;
   }

   for(i = 0; i < 4; i++)
   {
      if( fputc((FileSize >> (i * 8)) & 255, outfile) == EOF )
         return 1;
   }
   return 0;

} /* WriteHeader() */