content: {-# LANGUAGE BangPatterns, DeriveGeneric, DefaultSignatures #-}
  module Dust.Model.Huffman
  (
     HuffmanTree,
     encode,
     decode,
     bitpack,
     bitunpack,
     padToEight,
     codes,
     fileToTree,
     countsToTree
  )
  where
  
  import GHC.Generics
  import Data.Serialize hiding (encode, decode)
  
  import Data.Word
  import Data.Char (intToDigit)
  import Data.Bits ((.|.), shiftL, testBit)
  import Data.List (insertBy, foldl', sortBy, unfoldr)
  import Data.Maybe (fromJust)
  import Data.Ord (comparing)
  import qualified Data.ByteString           as S
  import qualified Data.ByteString.Lazy      as B
  import qualified Data.Map                  as M
  
  import Dust.Model.Stats (histogram)
  
  --------------------------------------------------
  
  data HuffmanTree a
    = LeafNode a Int
    | InternalNode Int (HuffmanTree a) (HuffmanTree a)
    deriving (Generic)
  instance (Serialize a) => Serialize (HuffmanTree a)
  
  -- build a multiline string representation of a huffman tree
  instance Show a => Show (HuffmanTree a) where
    show =
        go ""
      where
        spaces = map (const ' ')
        paren s = "(" ++ s ++ ")"
        go ss (LeafNode s o) = "--" ++ paren (show o) ++ show s ++ "\n"
        go ss (InternalNode o l r) =
            let root  = "--" ++ paren (show o) ++ "-+"
                ss' = ss ++ tail (spaces root)
                lbranch = go (ss' ++ "|") l
                rbranch = go (ss' ++ " ") r
            in root ++ lbranch
                    ++ ss' ++ "|\n"
                    ++ ss' ++ "`"
                    ++ rbranch
  
  frequency :: HuffmanTree a -> Int
  frequency (LeafNode     _ x  ) = x
  frequency (InternalNode x _ _) = x
  
  -- build a huffman tree bototm-up from a list of symbols sorted by frequency
  sortedHuffman :: [(a,Int)] -> HuffmanTree a
  sortedHuffman =
      -- first, convert each tuple into a Leaf, then combine
      combine . map toLeaf
    where
      -- repeatedly combine lowest frequency trees and reinsert the result into
      -- the frequency ordered list
      -- note: a priority queue could help
      combine [t] = t
      combine (ta:tb:ts) = combine . insertBy (comparing frequency) (merge ta tb) $ ts
      -- make an internal node from two trees. the frequency is the sum of the
      -- two trees frequencies
      merge ta tb = InternalNode (frequency ta + frequency tb) ta tb
      -- make a Leaf from a symbol,freq tuple
      toLeaf = uncurry LeafNode
  
  -- traverse the huffman tree generating a map from the symbol to its huffman
  -- tree path (where False is left, and True is right)
  codes :: Ord a => HuffmanTree a -> M.Map a [Bool]
  codes =
      M.fromList . go []
    where
      -- leaf nodes mark the end of a path to a symbol
      go p (LeafNode s _) = [(s,reverse p)]
      -- traverse both branches and accumulate a reverse path
      go p (InternalNode _ l r) = go (False:p) l ++ go (True:p) r
  
  -- from a table mapping symbols to their corresponding huffman tree bit paths,
  -- replace each instance of a symbol with its bit path
  encode :: Ord a => M.Map a [Bool] -> [a] -> [Bool]
  encode tbl =
      concatMap get
    where
      get x = fromJust (M.lookup x tbl)
  
  -- from a list of bits, navigate a given huffman tree and emit its decoded
  -- symbol when reaching a Leaf
  decode :: HuffmanTree a -> [Bool] -> [a]
  decode t0 xs0 =
      go t0 xs0
    where
      -- reached leaf, emit symbol
      go (LeafNode s _) bs = s : go t0 bs
      -- choose path based on bit
      go (InternalNode _ l r) (b:bs)
        | not b     = go l bs
        | otherwise = go r bs
      go _ [] = []
  
  --------------------------------------------------
  
  swap :: (a,b) -> (b,a)
  swap ~(a,b) = (b,a)
  
  showBits :: [Bool] -> String
  showBits = map (intToDigit . fromEnum)
  
  --------------------------------------------------
  
  bitpack :: [Bool] -> B.ByteString
  bitpack = B.pack . map packByte . takeWhile (not . null) . unfoldr (Just . splitAt 8)
      where
      packByte = foldl' (\i b -> (i `shiftL` 1) .|. (fromIntegral $ fromEnum b)) 0
  
  bitunpack :: S.ByteString -> [Bool]
  bitunpack = concatMap (\byte -> map (testBit byte) [7,6..0]) . S.unpack
  
  --------------------------------------------------
  
  padToEight :: [Bool] -> [Bool]
  padToEight bits =
      let len = length bits
          rem = len `mod` 8
          extra = 8-rem
          padding = replicate extra False
      in bits ++ padding
  
  fileToTree :: FilePath -> IO (HuffmanTree Word8)
  fileToTree path = do
      contents <- B.readFile path
      return $ bytesToTree $ B.unpack contents
  
  bytesToTree :: [Word8] -> (HuffmanTree Word8)
  bytesToTree text =
      let frequencies = histogram text
          sortedFrequencies = sortBy (comparing swap) frequencies
      in sortedHuffman sortedFrequencies
  
  countsToTree :: [(Word8, Int)] -> (HuffmanTree Word8)
  countsToTree frequencies =
      let sortedFrequencies = sortBy (comparing swap) frequencies
      in sortedHuffman sortedFrequencies
  