#!/usr/bin/env python
##                   (c) David MacKay - Free software. License: GPL
## For license statement see  http://www.gnu.org/copyleft/gpl.html
"""
This is a BLOCK compression algorithm
that uses the Huffman algorithm.

This simple block compressor assumes that the source file
is an exact multiple of the block length.
The encoding does not itself delimit the size of the file, so
the decoder needs to knows where the end of the compressed
file is. Thus we must either ensure the decoder knows
the original file's length, or we have to use a special
end-of-file character. A superior compressor would first
encode the source file's length at the head of the compressed
file; then the decoder would be able to stop at the right
point and correct any truncation errors. I'll fix this
in block2.py.

The package contains the following functions:

 findprobs(f=0.01,N=6):
    Find probabilities of all the events
    000000
    000001
     ...
    111111
    <-N ->

 Bencode(string,symbols,N):
    Reads in a string of 0s and 1s, forms blocks, and encodes with Huffman.

 Bdecode(string,root,N):
    Decode a binary string into blocks, then return appropriate 0s and 1s

 compress_it( inputfile, outputfile ):
    Make Huffman code, and compress

 uncompress_it( inputfile, outputfile ):
    Make Huffman code, and uncompress

 There are also three test functions.
 If block.py is run from a terminal, it invokes compress_it (using stdin)
 or uncompress_it (using stdin), respectively if there are zero arguments
 or one argument.

"""
## /home/mackay/python/compression/huffman/Huffman3.py
## This supplies the huffman algorithm, complete with encoders and decoders:

import sys, os
from typing import List

class InternalNode:
    name : str
    def __init__(self):
        self.leaf = 0
        self.child : List[InternalNode] = []
    def children(self,child0,child1):
        self.leaf = 0
        self.child.append(child0)
        self.child.append(child1)

class node:
    count : float
    index : int
    name  : str
    internalnode : InternalNode
    def __init__(self, count, index , name=""):
        self.count = float(count)
        self.index = index
        self.name  = name ## optional argument
        if self.name=="" : self.name = '_'+str(index)
        self.word = "" ## codeword will go here
        self.isinternal = 0
    # def __cmp__(self, other):
    #     return cmp(self.count, other.count)
    def __lt__(self, other):
        return self.count < other.count
    def report(self):
        if (self.index == 1 ) :
            print('#Symbol\tCount\tCodeword')
        print('%s\t(%f)\t%s' % (self.name,self.count,self.word))
    def associate(self,internalnode):
        self.internalnode = internalnode
        self.internalnode.leaf = 1
        self.internalnode.name = self.name

def find_idx(seq, index):
    for item in seq:
        if item.index == index:
            return item

def find_name(seq, name):
    for item in seq:
        if item.name == name:
            return item
    return node(0, -1)

def iterate (c : List[node]) -> InternalNode:
    """
    Run the Huffman algorithm on the list of "nodes" c.
    The list of nodes c is destroyed as we go, then recreated.
    Codewords 'co.word' are assigned to each node during the recreation of the list.
    The order of the recreated list may well be different.
    Use the list c for encoding.

    The root of a new tree of "internalnodes" is returned.
    This root should be used when decoding.

    >>> c = [ node(0.5,1,'a'),  \
              node(0.25,2,'b'), \
              node(0.125,3,'c'),\
              node(0.125,4,'d') ]   # my doctest query has been resolved
    >>> root = iterate(c)           # "iterate(c)" returns a node, not nothing, and doctest cares!
    >>> reportcode(c)               # doctest: +NORMALIZE_WHITESPACE, +ELLIPSIS
    #Symbol   Count	Codeword
    a         (0.5)	1
    b         (0.25)	01
    c         (0.12)	000
    d         (0.12)	001
    """
    root : InternalNode
    if ( len(c) > 1 ) :
        c.sort() ## sort the nodes by count, using the __cmp__ function defined in the node class
        deletednode = c[0] ## keep copy of smallest node so that we can put it back in later
        second = c[1].index ## index of second smallest node
        # MERGE THE BOTTOM TWO
        c[1].count += c[0].count  ##  this merged node retains the name of the bigger leaf.
        c.pop(0)

        root = iterate ( c )

        ## Fill in the new information in the ENCODING list, c
        ## find the codeword that has been split/joined at this step
        co = find_idx(c, second)
#        co = find( lambda p: p.index == second , c )
        deletednode.word = co.word+'0'
        c.append( deletednode )  ## smaller one gets 0
        co.word += '1'
        co.count -= deletednode.count   ## restore correct count

        ## make the new branches in the DECODING tree
        newnode0 = InternalNode()
        newnode1 = InternalNode()
        treenode = co.internalnode # find the node that got two children
        treenode.children(newnode0,newnode1)
        deletednode.associate(newnode0)
        co.associate(newnode1)
    else :
        c[0].word = ""
        root = InternalNode()
        c[0].associate(root)
    return root

def encode(sourcelist,code):
    """
    Takes a list of source symbols. Returns a binary string.
    """
    answer = ""
    for s in sourcelist:
        co = find_name(code, s)
#        co = find(lambda p: p.name == s, code)
        if co.index == -1:
            print("Warning: symbol " + s + " has no encoding!")
        else:
            answer = answer + co.word
    return answer

def decode(string,root):
    """
    Decodes a binary string using the Huffman tree accessed via root
    """
    ## split the string into a list
    ## then copy the elements of the list one by one.
    answer : List[str] = []
    clist = list( string )
    ## start from root
    currentnode = root
    for c in clist:
        if ( c=='\n' ):  continue ## special case for newline characters
        assert ( c == '0' )or( c == '1')
        currentnode = currentnode.child[int(c)]
        if currentnode.leaf != 0:
            answer.append( str(currentnode.name) )
            currentnode = root
    #assert (currentnode == root) ## if this is not true then we have run out of characters and are half-way through a codeword
    return answer

## alternate way of calling huffman with a list of counts ## for use as package by other programs ######
## two example ways of using it:
#>>> from Huffman3 import *
#>>> huffman([1, 2, 3, 4],1)
#>>> (c,root) = huffman([1, 2, 3, 4])

## end ##########################################################################

def makenodes(probs):
    """
    Creates a list of nodes ready for the Huffman algorithm.
    Each node will receive a codeword when Huffman algorithm "iterate" runs.

    probs should be a list of pairs('<symbol>', <value>).

    >>> probs=[('a',0.5), ('b',0.25), ('c',0.125), ('d',0.125)]
    >>> symbols = makenodes(probs)
    >>> root = iterate(symbols)
    >>> zipped = encode(['a','a','b','a','c','b','c','d'], symbols)
    >>> print zipped
    1101100001000001
    >>> print decode( zipped, root )
    ['a', 'a', 'b', 'a', 'c', 'b', 'c', 'd']

    See also the file Example.py for a python program that uses this package.
    """
    m=0
    c : List[node]=[]
    for p in probs:
        m += 1
        c.append( node( p[1], m, p[0] ) )
    return c

def dec_to_bin( n , digits ):
    """ n is the number to convert to binary;  digits is the number of bits you want
    Always prints full number of digits
    >>> print dec_to_bin( 17 , 9)
    000010001
    >>> print dec_to_bin( 17 , 5)
    10001

    Will behead the standard binary number if requested
    >>> print dec_to_bin( 17 , 4)
    0001
    """
    if(n<0) :
        sys.stderr.write( "warning, negative n not expected\n")
    i=digits-1
    ans=""
    while i>=0 :
        b = (((1<<i)&n)>0)
        i -= 1
        ans = ans + str(int(b))
    return ans

verbose=0

def weight(string):
    """
    ## returns number of 0s and number of 1s in the string
    >>> print weight("00011")
    (3, 2)
    """
    w0=0;w1=0
    for c in list(string):
        if(c=='0'):
            w0+=1
        elif(c=='1'):
            w1+=1
    return (w0,w1)


def findprobs(f=0.01,nn=6):
    """ Find probabilities of all the events
    000000
    000001
     ...
    111111
    <-nn ->
    >>> print findprobs(0.1,3)              # doctest:+ELLIPSIS
    [('000', 0.7...),..., ('111', 0.001...)]
    """
    answer : List[Tuple[str, float]] = []
    for n in range(int(2**nn)):
        s = dec_to_bin(n,nn)
        (w0,w1) = weight(s)
        if verbose and 0 :
            print(s + ' ' + w0 + ' ' + w1)
        answer.append( (s, f**w1 * (1-f)**w0 ) )
    assert ( len(answer) == 2**nn )
    return answer

def Bencode(string,symbols,n):
    """
    Reads in a string of 0s and 1s.
    Creates a list of blocks of size n.
    Sends this list to the general-purpose Huffman encoder
    defined by the nodes in the list "symbols".
    """
    blocks : List[str] = []
    chars = list(string)

    s=""
    for c in chars:
        s = s+c
        if (len(s)>=n):
            blocks.append( s )
            s = ""
    if (len(s)>0):
        print("warning, padding last block with 0s")
        while (len(s)<n):
            s = s+'0'
        blocks.append( s )

    if verbose:
        print("blocks to be encoded:")
        print(blocks)
    zipped = encode( blocks , symbols )
    return zipped

def Bdecode(string,root,n):
    """
    Decode a binary string into blocks.
    """
    answer = decode( string, root )
    if verbose:
        print("blocks from decoder:")
        print(answer)
    output = "".join( answer )
    ## this assumes that the source file was an exact multiple of the blocklength
    return output

def easytest():
    """
    Tests block code with N=3, f=0.01 on a tiny example.
    >>> easytest()                 # doctest:+NORMALIZE_WHITESPACE
    #Symbol	Count	        Codeword
    000	        (0.97)	        1
    001	        (0.0098)	001
    010	        (0.0098)	010
    011	        (9.9e-05)	00001
    100	        (0.0098)	011
    101	        (9.9e-05)	00010
    110	        (9.9e-05)	00011
    111	        (1e-06) 	00000
    zipped  = 1001010000010110111
    decoded = ['000', '001', '010', '011', '100', '100', '000']
    OK!
    """
    n=3
    f=0.01
    probs = findprobs(f,n)
#    if len(probs) > 999 :
#        sys.setrecursionlimit( len(probs)+100 )
    symbols = makenodes(probs) # makenodes is defined at the bottom of Huffman3 package
    root = iterate(symbols) # make huffman code and put it into the symbols' nodes, and return the root of the decoding tree

    symbols = sorted(symbols, key = lambda x: x.index) # sort by index
    for co in symbols :                                # and write the answer
        co.report()

    source = ['000','001','010','011','100','100','000']
    zipped = encode(source, symbols)
    print("zipped  = " + zipped)
    answer = decode( zipped, root )
    print("decoded = ", end = ''); print(answer)
    if ( source != answer ):
        print("ERROR")
    else:
        print("OK!")

f=0.01; n=12   #  1244 bits if n==12
f=0.01; n=5   #  2266  bits if n==5
f=0.01; n=10   #  1379 bits if n==10

def compress_it( inputfile, outputfile ):
    """
    Make Huffman code for blocks, and
    Compress from file (possibly stdin).
    """
    probs = findprobs(f,n)
    symbols = makenodes(probs)
#    if len(probs) > 999 :
#        sys.setrecursionlimit( len(probs)+100 )
    root = iterate(symbols) # make huffman code and put it into the symbols' nodes, and return the root of the decoding tree

    string = inputfile.read()
    outputfile.write( Bencode(string, symbols, n) )

def uncompress_it( inputfile, outputfile ):
    """
    Make Huffman code for blocks, and
    UNCompress from file (possibly stdin).
    """
    probs = findprobs(f,n)
#    if len(probs) > 999 :
#        sys.setrecursionlimit( len(probs)+100 )
    symbols = makenodes(probs)
    root = iterate(symbols) # make huffman code and put it into the symbols' nodes, and return the root of the decoding tree

    string = inputfile.read()
    outputfile.write( Bdecode(string, root, n) )

def hardertest():
    print("Reading the BentCoinFile")
    inputfile = open( "testdata/BentCoinFile" , "r" )
    outputfile = open( "tmp.zip" , "w" )
    print("Compressing to tmp.zip")
    compress_it(inputfile, outputfile)
    outputfile.close();     inputfile.close()
#    print "DONE compressing"

    inputfile = open( "tmp.zip" , "r" )
    outputfile = open( "tmp2" , "w" )
    print("Uncompressing to tmp2")
    uncompress_it(inputfile, outputfile)
    outputfile.close();     inputfile.close()
#    print "DONE uncompressing"

    print("Checking for differences...")
    print((open('testdata/BentCoinFile').read() == open('tmp2').read()))#os.system( "diff testdata/BentCoinFile tmp2" )
    #os.system( "wc tmp.zip testdata/BentCoinFile tmp2" )

def test():
    easytest()
    hardertest()

if __name__ == '__main__':
    #sys.setrecursionlimit( 10000 )
    test()

