using System;
using System.Collections;
using System.Collections.Generic;
using System.Text;

namespace QiHe.CodeLib.Compress
{
    public partial class HuffmanTree
    {
        public int[] CodeLengths;
        int[] Codes;
        BinaryTree Tree;
        public HuffmanTree(int[] codeLengths)
        {
            CodeLengths = codeLengths;
            Codes = DeflateHuffmanCoding.AssignCodes(codeLengths);
        }

        public static HuffmanTree FromCodeLengths(int[] codeLengths)
        {
            HuffmanTree tree = new HuffmanTree(codeLengths);
            tree.BuildTree();
            return tree;
        }

        public int ReadSymbol(BitStream input)
        {
            long pos = input.Position;
            int bits = input.bitsRead;
            BinaryTreeNode node = Tree.Root;
            while (node is BinaryTreeInternalNode)
            {
                int bit = input.ReadBit();
                if (bit == -1)
                {
                    throw new Exception("data corupted");
                }
                node = ((BinaryTreeInternalNode)node).ChildNodes[bit];
            }
            return ((BinaryTreeLeafNode)node).Symbol;
        }

        public void WriteSymbol(BitStream output, int symbol)
        {
            output.WriteBitsBigEndian(Codes[symbol], CodeLengths[symbol]);
        }

        public string ListSymbolCodes()
        {
            StringBuilder text = new StringBuilder();
            foreach (Pair<int, string> pair in Tree.GetSymbolCodePairs(true))
            {
                text.AppendFormat("{0}\t{1}\r\n", pair.Left, pair.Right);
            }
            return text.ToString();
        }

        private void BuildTree()
        {
            BinaryTreeInternalNode root = new BinaryTreeInternalNode();
            for (int symbol = 0; symbol < Codes.Length; symbol++)
            {
                int length = CodeLengths[symbol];
                if (length == 0) continue; //!important
                int code = Codes[symbol];
                BitArray bits = new BitArray(BitConverter.GetBytes(code));
                BinaryTreeInternalNode parent = root;
                for (int i = 0; i < length - 1; i++)
                {
                    parent = CreateInternalNode(parent, Convert.ToInt32(bits[length - 1 - i]));
                }
                BinaryTreeLeafNode leaf = new BinaryTreeLeafNode(symbol);
                parent.ChildNodes[Convert.ToInt32(bits[0])] = leaf;
            }
            Tree = new BinaryTree(root);
        }

        static BinaryTreeInternalNode CreateInternalNode(BinaryTreeInternalNode parent, int index)
        {
            if (parent.ChildNodes[index] == null)
            {
                parent.ChildNodes[index] = new BinaryTreeInternalNode();
            }
            if (parent.ChildNodes[index] is BinaryTreeInternalNode)
            {
                return parent.ChildNodes[index] as BinaryTreeInternalNode;
            }
            else
            {
                throw new Exception("Inconsistent Huffman tree structure.");
            }
        }
    }
}