USING: accessors arrays assocs bit-arrays               
combinators.short-circuit lists fry kernel make math math.order
sequences sorting summary ;                                    
IN: huffman-coding                                             

TUPLE: huffman-code tree-root dictionary ;

<PRIVATE

TUPLE: huffman-leaf symbol probability ;
TUPLE: huffman-node probability left right ;

C: <huffman-leaf> huffman-leaf
C: <huffman-node> huffman-node

: combine-nodes ( a b -- c )
    [ [ probability>> ] bi@ + ] 2keep
    <huffman-node> ;                 

: (huffman-sorted-insert) ( list node -- )
    2dup [ cadr probability>> ] [ probability>> ] bi* >
    [ over cdr cons >>cdr drop ]                       
    [ [ cdr ] dip (huffman-sorted-insert) ] if ;       

: huffman-sorted-insert ( list node -- list )
    2dup {                                   
        [ drop nil? ]                        
        [ [ car probability>> ] [ probability>> ] bi* >= ]
    } 2||                                                 
    [ swap cons ] [ dupd (huffman-sorted-insert) ] if ; inline

: build-huffman-tree ( list -- tree )
    dup cdr nil? not [               
        [ cdr cdr ] [ 2car ] bi      
        combine-nodes huffman-sorted-insert
        build-huffman-tree                 
    ] [ car ] if ;                         

GENERIC# (build-dictionary) 1 ( internal-tree accum -- )

M: huffman-leaf (build-dictionary)
    [ symbol>> ] dip >bit-array 2array , ;

M: huffman-node (build-dictionary)
    [ [ left>> ] [ right>> ] bi ] dip
    dup clone                        
    [ dup t swap push ] [ dup f swap push ] bi*
    swapd [ (build-dictionary) ] 2bi@ ;        

: build-dictionary ( internal-tree -- dictionary )
    [ V{ } (build-dictionary) ] { } make H{ } assoc-like ;

PRIVATE>

: <huffman-code> ( assoc -- huffman-code )
    [ <huffman-leaf> ] { } assoc>map      
    [ [ probability>> ] bi@ <=> ] sort seq>list
    build-huffman-tree dup build-dictionary    
    huffman-code boa ;                         

<PRIVATE

ERROR: illegal-huffman-symbol symbol ;

M: illegal-huffman-symbol summary
    drop "You tried to encode a symbol which is not in the huffman-code" ;

PRIVATE>

: huffman-encode ( huffman-code seq -- code )
    swap dictionary>> '[                     
        dup _ at* [ nip ]                    
        [ drop illegal-huffman-symbol ] if   
    ] { } map-as concat ;                    

<PRIVATE

TUPLE: huffman-decoder root current ;

: <huffman-decoder> ( huffman-code -- huffman-decoder )
    tree-root>> dup huffman-decoder boa ;              

GENERIC: update-decoder ( decoder tree -- decoder )

M: huffman-node update-decoder
    >>current ;

M: huffman-leaf update-decoder
    symbol>> , dup root>> >>current ;

: decode-step ( huffman-decoder bit -- huffman-decoder )
    [ dup current>> ] dip
    [ left>> ] [ right>> ] if
    update-decoder ;

PRIVATE>

: huffman-decode-as ( huffman-code code exemplar -- seq )
    [
        [ <huffman-decoder> ] dip
        [ decode-step ]
        [ each ]
    ] dip make nip ;

: huffman-decode ( huffman-code code -- array )
    { } huffman-decode-as ; inline

<PRIVATE

: get-statistics ( seq -- hash )
    H{ } clone [ '[ _ inc-at ] each ] keep ;

PRIVATE>

: huffman-count-and-encode ( seq -- huffman-code code )
    [ get-statistics <huffman-code> dup ] keep
    huffman-encode ;