#!/usr/bin/runghc

type Symbol = Char

type Weight = Float

data Node
  = Leaf Symbol Weight
  | Interior Node Node [Symbol] Weight
  deriving (Show)

getSymbols :: Node -> [Symbol]
getSymbols (Leaf s _) = [s]
getSymbols (Interior _ _ syms _) = syms

getWeight :: Node -> Weight
getWeight (Leaf _ w) = w
getWeight (Interior _ _ _ w) = w

orderedInsert :: Node -> [Node] -> [Node]
orderedInsert n [] = [n]
orderedInsert n all@(first : rest)
  | getWeight n <= getWeight first = n : all
  | otherwise = first : orderedInsert n rest

makeTree :: [Node] -> Node
makeTree [n] = n
makeTree (first : second : rest) =
  let n =
        Interior
          first
          second
          (getSymbols first ++ getSymbols second)
          (getWeight first + getWeight second)
   in makeTree (orderedInsert n rest)

-- These are derived from a large standard corpus of text.
commonTree =
  makeTree
    [ Leaf ' ' 30.2,
      Leaf 'e' 12.702,
      Leaf 't' 9.056,
      Leaf 'a' 8.167,
      Leaf 'o' 7.507,
      Leaf 'i' 6.966,
      Leaf 'n' 6.749,
      Leaf 's' 6.327,
      Leaf 'h' 6.094,
      Leaf 'r' 5.987,
      Leaf 'd' 4.253,
      Leaf 'l' 4.025,
      Leaf 'c' 2.782,
      Leaf 'u' 2.758,
      Leaf 'm' 2.406,
      Leaf 'w' 2.360,
      Leaf 'f' 2.228,
      Leaf 'g' 2.015,
      Leaf 'y' 1.974,
      Leaf 'p' 1.929,
      Leaf 'b' 1.492,
      Leaf 'v' 0.978,
      Leaf 'k' 0.772,
      Leaf 'j' 0.153,
      Leaf 'x' 0.150,
      Leaf 'q' 0.095,
      Leaf 'z' 0.071
    ]

encodeRec :: Node -> [Symbol] -> [Bool]
encodeRec _ [] = []
encodeRec (Leaf s w) (sym : symbols) = encodeRec commonTree symbols
encodeRec (Interior l r _ _) (sym : symbols)
  | elem sym (getSymbols l) = False : encodeRec l (sym : symbols)
  | otherwise = True : encodeRec r (sym : symbols)

encode = encodeRec commonTree

decodeRec :: Node -> [Bool] -> [Symbol]
decodeRec (Leaf s w) bits = s : decodeRec commonTree bits
decodeRec (Interior l _ _ _) (False : bits) = decodeRec l bits
decodeRec (Interior _ r _ _) (True : bits) = decodeRec r bits
decodeRec _ [] = []

decode = decodeRec commonTree

main :: IO ()
main = do
  putStrLn ("This is a generalized tree:\n" ++ show commonTree)

  let message = "this is a small message"
  putStrLn ("\nMessage before encoding was:\n" ++ message)
  putStrLn ("\nLength of message before encoding, times 8 was:\n" ++ show (8 * length message))

  let encoded = encode message
  putStrLn ("\nEncoded message was:\n" ++ show encoded)
  putStrLn ("\nLength of encoded message was:\n" ++ show (length encoded))

  let decoded = decode encoded
  putStrLn ("\nThe decoded message was:\n" ++ decoded)