{-# LINE 159 "asymm.lhs" #-}
{- 

Code extracted from the paper

  "Coding with Asymmetric Numeral Systems"
  by Jeremy Gibbons

to appear in Maths of Program Construction, Oct 2019.  

 -}

import Data.Ratio (Rational, (%))
import Data.List (unfoldr)
import Test.QuickCheck hiding (scale, total)
import Data.Char (ord, chr)
import Data.Maybe (fromJust)
{-# LINE 257 "asymm.lhs" #-}
huffman :: [Symbol] -> [Bool]
huffman = concatMap code where
  code 'a'  = [False,False]
  code 'b'  = [False,True]
  code 'c'  = [True]
{-# LINE 520 "asymm.lhs" #-}
fmapList :: (b -> c) -> Maybe (a,b) -> Maybe (a,c)
fmapList h (Just (a,b))  = Just (a, h b)
fmapList g Nothing       = Nothing
{-# LINE 679 "asymm.lhs" #-}
stream :: (b -> Maybe (c,b)) -> (b -> a -> b) -> b -> [a] -> [c]
stream g f b x = case g b of 
    Just (c,b')  -> c : stream g f b' x
    Nothing      -> case x of 
      a:x'  -> stream g f (f b a) x'  
      []    -> []
{-# LINE 695 "asymm.lhs" #-}
fstream :: (b -> Maybe (c,b)) -> (b -> [c]) -> (b -> a -> b) -> b -> [a] -> [c]
fstream g h f b x = case g b of 
    Just (c,b')  -> c : fstream g h f b' x
    Nothing      -> case x of 
      a:x'  -> fstream g h f (f b a) x' 
      []    -> h b
{-# LINE 711 "asymm.lhs" #-}
apo :: (b -> Maybe (c,b)) -> (b -> [c]) -> b -> [c]
apo g h b = case g b of
  Just (c, b')  -> apo g h b'
  Nothing       -> h b
{-# LINE 728 "asymm.lhs" #-}
guard :: (b -> Bool) -> (b -> Maybe (c,b)) -> (b -> Maybe (c,b)) 
guard p g x = if  p x  then  g x  else  Nothing
{-# LINE 744 "asymm.lhs" #-}
type Symbol = Char

counts = [ ('a',2), ('b',3), ('c',5) ]

intFromSymbol :: Symbol -> Int
intFromSymbol s = ord s - ord 'a'
intToSymbol :: Int -> Symbol
intToSymbol n = chr (ord 'a' + n)
fromSymbol :: Symbol -> Integer
fromSymbol = toInteger . intFromSymbol
toSymbol :: Integer -> Symbol
toSymbol = intToSymbol . fromInteger

newtype CheckSymbol = S { unS :: Symbol } deriving Eq
instance Show CheckSymbol where
  show (S s) = show s
  showList = showList . map unS
instance Arbitrary CheckSymbol where
  arbitrary = frequency [ (fromIntegral c, return (S s)) | (s,c) <- counts ]
instance Enum CheckSymbol where
  fromEnum = intFromSymbol . unS
  toEnum = S . intToSymbol

unit :: Interval
unit = (0,1)
{-# LINE 781 "asymm.lhs" #-}
type Interval = (Rational,Rational)
{-# LINE 788 "asymm.lhs" #-}
counts :: Integral n => [(Symbol,n)]
{-# LINE 795 "asymm.lhs" #-}
symbols :: [Symbol]
symbols = map fst counts

total :: Integral n => n
total = sum (map snd counts)

freqs :: [(Symbol,Rational)]
freqs = [ (s,c % total) | (s,c) <- counts ]

freq :: Symbol -> Rational
freq = justLookup freqs

fcumuls :: [(Symbol,Rational)]
fcumuls = zip symbols (scanl (+) 0 (map snd freqs))

fcumul :: Symbol -> Rational
fcumul = justLookup fcumuls

encodeSym :: Symbol -> Interval 
encodeSym s  = (l, l + freq s) where l = fcumul s

decodeSym :: Rational -> Symbol
decodeSym x  = last [ s | (s,y) <- fcumuls, y <= x ]
{-# LINE 821 "asymm.lhs" #-}
justLookup :: Eq a => [(a,b)] -> a -> b
justLookup abs a = fromJust (lookup a abs)
{-# LINE 849 "asymm.lhs" #-}
weight, scale  :: Interval -> Rational -> Rational
weight (l,r) x  = l + (r-l) * x 
scale (l,r) y   = (y-l)/(r-l) 

narrow         :: Interval -> Interval -> Interval
narrow i (p,q)  = (weight i p, weight i q) 
{-# LINE 870 "asymm.lhs" #-}
encode1  :: [Symbol] -> Rational 
encode1  = fst . foldl estep1 unit

estep1 :: Interval -> Symbol -> Interval 
estep1 i s = narrow i (encodeSym s)
  
decode1 :: Rational -> [Symbol] 
decode1 = unfoldr dstep1

dstep1 :: Rational -> Maybe (Symbol, Rational) 
dstep1 x =  let s = decodeSym x in Just (s, scale (encodeSym s) x)
{-# LINE 910 "asymm.lhs" #-}
prop_EncodeDecode1 :: [CheckSymbol] -> Bool
prop_EncodeDecode1 xs = take (length t) (decode1 (encode1 t)) == t
  where t = map unS xs

encode1'  :: [Symbol] -> Interval
encode1'  = foldl estep1 unit  

pick :: Interval -> Rational
pick i@(l,r)
    | 1/2 < l    = (1 + pick (narrow (-1,1) i)) / 2
    | r <= 1/2   = pick (narrow (0,2) i) / 2
    | otherwise  = 1/2

prop_EncodeDecode1' :: [CheckSymbol] -> Bool
prop_EncodeDecode1' xs = take (length t) (decode1 (pick (encode1' t))) == t
  where t = map unS xs

bits :: Interval -> [Bool]
bits i@(l,r)
    | 1/2 < l    = True : bits (narrow (-1,1) i)
    | r <= 1/2   = False : bits (narrow (0,2) i) 
    | otherwise  = []

prop_EncodeHuffman :: [CheckSymbol] -> Bool
prop_EncodeHuffman xs = length (huffman t) - 1 <= length (bits (encode1' t))
  where t = map unS xs
{-# LINE 984 "asymm.lhs" #-}
encode2  :: [Symbol] -> Rational 
encode2  = foldr estep2 0 

estep2 :: Symbol -> Rational -> Rational
estep2 s x = weight (encodeSym s) x
{-# LINE 1073 "asymm.lhs" #-}
prop_EncodeDecode2 :: [CheckSymbol] -> Bool
prop_EncodeDecode2 xs = take (length t) (decode1 (encode2 t)) == t
  where t = map unS xs
{-# LINE 1092 "asymm.lhs" #-}
count   ::  Integral n => Symbol -> n
count   =   justLookup counts

cumul   ::  Integral n => Symbol -> n
cumul   =   justLookup cumuls

cumuls  ::  Integral n => [(Symbol, n)]
cumuls  =   zip symbols (scanl (+) 0 (map snd counts))

find    ::  Integral n => n -> Symbol
find x  =   last [ s | (s,y) <- cumuls, y <= x ] --  assuming |x < total|
{-# LINE 1190 "asymm.lhs" #-}
encode3 :: [Symbol] -> Integer 
encode3  = foldr estep3 0 

estep3 :: Integral n => Symbol -> n -> n
estep3 s x = let (q,r) = x `divMod` count s in q * total + cumul s + r
  
decode3  :: Integer -> [Symbol] 
decode3 = unfoldr dstep3 

dstep3 :: Integral n => n -> Maybe (Symbol, n)
dstep3 x =  let  (q,r) = x `divMod` total 
                 s = find r 
            in Just (s, count s * q + r - cumul s)
{-# LINE 1265 "asymm.lhs" #-}
prop_EncodeDecode3 :: [CheckSymbol] -> Bool
prop_EncodeDecode3 xs = take (length t) (decode3 (encode3 t)) == t
  where t = map unS xs
{-# LINE 1298 "asymm.lhs" #-}
funnyHuffman :: [Symbol] -> [Int]
funnyHuffman = prune . foldr hstep [0,0] where
  hstep s xs = let n = 3 - length (code s)
                   (ys,zs) = splitAt (length xs - n) xs 
               in ys ++ code s ++ zs
  code 'a'  = [0]
  code 'b'  = [1,0]
  code 'c'  = [1,1,0]
  code 'd'  = [1,1,1]
  prune [0] = [0]
  prune (0:xs) = prune xs
  prune (1:xs) = 1:xs

encode3' :: [Symbol] -> Integer --  ANS with Huffman-friendly counts
encode3'  = foldr estep3 0 where
  counts   = [ ('a',4), ('b',2), ('c',1), ('d',1) ]
  symbols  = map fst counts
  total    = sum (map snd counts)
  count    = justLookup counts
  cumul    = justLookup (zip symbols (scanl (+) 0 (map snd counts)))
  estep3 s x = let (q,r) = x `divMod` count s in q * total + cumul s + r
  
toBinary :: Integer -> [Int]
toBinary n | n<2 = [fromInteger n]
toBinary n = toBinary m ++ [fromInteger b] where (m,b) = n `divMod` 2

prop_Huffman :: [CheckSymbol] -> Bool  --  this only checks alphabet a,b,c
prop_Huffman xs = funnyHuffman t == toBinary (encode3' t)
  where t = map unS xs
{-# LINE 1341 "asymm.lhs" #-}
encode4 :: [Symbol] -> Integer 
encode4  = foldr estep3 l

decode4  :: Integer -> [Symbol] 
decode4 = unfoldr dstep4 

dstep4 :: Integer -> Maybe (Symbol, Integer) 
dstep4 x = if  x==l  then  Nothing  else  dstep3 x
{-# LINE 1358 "asymm.lhs" #-}
prop_EncodeDecode4 :: [CheckSymbol] -> Bool
prop_EncodeDecode4 xs = decode4 (encode4 t) == t
  where t = map unS xs
{-# LINE 1370 "asymm.lhs" #-}
b, l, u :: Integral n => n
b = 10
l = 100
u = l * b

inject :: Integral n => n -> n -> n
extract :: Integral n => n -> (n,n)
{-# LINE 1391 "asymm.lhs" #-}
type  Number = (Int, [Int])
{-# LINE 1400 "asymm.lhs" #-}
abstract :: Number -> Integer
abstract (w,ys) = foldl inject (fromIntegral w) (map fromIntegral ys)
{-# LINE 1411 "asymm.lhs" #-}
inject w y   = w * b + y 
{-# LINE 1421 "asymm.lhs" #-}
extract w    = w `divMod` b
{-# LINE 1509 "asymm.lhs" #-}
econsume5 :: [Symbol] -> Number 
econsume5  = foldr estep5 (l,[])

estep5 :: Symbol -> Number -> Number 
estep5 s (w,ys) = let (w',ys') = enorm5 s (w,ys) in (estep3 s w',ys')

enorm5 :: Symbol -> Number -> Number 
enorm5 s (w,ys) =  if    estep3 s w < u  
                   then  (w,ys) 
                   else  let (q,r) = extract w  in  enorm5 s (q,r:ys)
{-# LINE 1589 "asymm.lhs" #-}
dproduce5 :: Number -> [Symbol] 
dproduce5 = unfoldr dstep5

dstep5 :: Number -> Maybe (Symbol, Number) 
dstep5 (w,ys) =  let  Just (s, w')  = dstep3 w
                      (w'',ys'')    = dnorm5 (w',ys)  
                 in if  w'' >= l  then  Just (s,(w'',ys''))  else  Nothing

dnorm5 :: Number -> Number 
dnorm5 (w,y:ys)  = if  w < l  then  dnorm5  (inject w y, ys)  else  (w,y:ys) 
dnorm5 (w,[])    = (w,[])
{-# LINE 1751 "asymm.lhs" #-}
prop_EconsumeDproduce5 :: [CheckSymbol] -> Bool
prop_EconsumeDproduce5 xs = dproduce5 (econsume5 t) == t
  where t = map unS xs
{-# LINE 1769 "asymm.lhs" #-}
eflush5 :: Number -> [Int] 
eflush5 (0,ys)  = ys 
eflush5 (w,ys)  = let  (w',y) = extract w  in  eflush5 (w',y:ys)
{-# LINE 1776 "asymm.lhs" #-}
encode5 :: [Symbol] -> [Int] 
encode5  = eflush5 . econsume5
{-# LINE 1782 "asymm.lhs" #-}
dstart5 :: [Int] -> Number 
dstart5 ys = dnorm5 (0,ys)
{-# LINE 1788 "asymm.lhs" #-}
decode5 :: [Int] -> [Symbol] 
decode5  = dproduce5 . dstart5
{-# LINE 1832 "asymm.lhs" #-}
prop_EncodeDecode5 :: [CheckSymbol] -> Bool
prop_EncodeDecode5 xs = decode5 (encode5 t) == t
  where t = map unS xs
{-# LINE 1861 "asymm.lhs" #-}
edeal6 :: [alpha] -> [alpha] 
edeal6 = unfoldr unsnoc  where 
  unsnoc []  = Nothing 
  unsnoc ys  = let  (ys',y) = (init ys, last ys)  in  Just (y, ys')
{-# LINE 1954 "asymm.lhs" #-}
efstep6 :: Number -> Maybe (Int, Number)
efstep6 (0,[])  = Nothing 
efstep6 (w,[])  = let  (q,r) = extract w  in  Just (r,(q,[])) 
efstep6 (w,ys)  = let  (ys',y) = (init ys, last ys)  in  Just (y, (w,ys'))
{-# LINE 1961 "asymm.lhs" #-}
eflush6 :: Number -> [Int] 
eflush6 = reverse . unfoldr efstep6
{-# LINE 1971 "asymm.lhs" #-}
encode6 :: [Symbol] -> [Int] 
encode6  = reverse . unfoldr efstep6 . foldl (flip estep5) (l,[]) . reverse
{-# LINE 1977 "asymm.lhs" #-}
prop_EncodeDecode6 :: [CheckSymbol] -> Bool
prop_EncodeDecode6 xs = decode5 (encode6 t) == t
  where t = map unS xs
{-# LINE 2003 "asymm.lhs" #-}
efstep7 :: Number -> Maybe (Int, Number) 
efstep7 = guard (not . null . snd) efstep6
{-# LINE 2080 "asymm.lhs" #-}
encode7 :: [Symbol] -> [Int] 
encode7  = reverse . fstream efstep7 (unfoldr efstep6) (flip estep5) (l,[]) . reverse
{-# LINE 2089 "asymm.lhs" #-}
prop_EncodeDecode7 :: [CheckSymbol] -> Bool
prop_EncodeDecode7 xs = decode5 (encode7 t) == t
  where t = map unS xs
{-# LINE 2110 "asymm.lhs" #-}
dsstep8 (w,[]) y  = if  w < l  then  (inject w y, [])  else  (w, [y]) 
dsstep8 (w,ys) y  = (w, ys ++ [y])
{-# LINE 2159 "asymm.lhs" #-}
decode8 :: [Int] -> [Symbol] 
decode8  = stream dstep5 dsstep8 (0,[])
{-# LINE 2166 "asymm.lhs" #-}
prop_EncodeDecode8 :: [CheckSymbol] -> Bool
prop_EncodeDecode8 xs = decode8 (encode7 t) == t
  where t = map unS xs
{-# LINE 2178 "asymm.lhs" #-}
encode9 :: [Symbol] -> [Int] 
encode9 = reverse . fstream g' (unfoldr g) f (l,[]) . reverse where 
  f (w,ys) s  =  if    w < b * (l `div` total) * count s 
                 then  let  (q,r) = w `divMod` count s  in  (q * total + cumul s + r,ys) 
                 else  let  (q,r) = extract w  in  f (q, r : ys) s
  g (0,[])    = Nothing 
  g (w,[])    = let  (q,r) = extract w  in  Just (r,(q,[])) 
  g (w,ys)    = let  (ys',y) = (init ys, last ys)  in  Just (y, (w,ys')) 
  g' (w,ys)   = if  null ys  then  Nothing  else  g (w,ys)
{-# LINE 2190 "asymm.lhs" #-}
decode9 :: [Int] -> [Symbol] 
decode9 = stream g f (0,[]) where 
  g (w,ys)    =  let   (q,r)  = w `divMod` total 
                       s  = find r 
                       (w'',ys'')  = n (count s * q + r - cumul s,ys) 
                 in if  w'' >= l  then  Just (s,(w'',ys''))  else  Nothing
  n (w,ys)    =  if w>= l || null ys  then  (w,ys)  else 
                 let  (y:ys') = ys  in  n (inject w y, ys')

  f (w,[]) y  = n (w,[y]) 
  f (w,ys) y  = (w, ys ++ [y])
{-# LINE 2210 "asymm.lhs" #-}
prop_EncodeDecode9 :: [CheckSymbol] -> Bool
prop_EncodeDecode9 xs = decode9 (encode9 t) == t
  where t = map unS xs
{-# LINE 2297 "asymm.lhs" #-}
encode10 :: [Symbol] -> [Int] 
encode10 = reverse . e1 l . reverse where
  e1 w (s:ss) =  let  (q,r)  = w `divMod` count s
                      w'     = q * total + cumul s + r in
                 if w' < u  then  e1 w' ss
                            else  let (q',r') = w `divMod` b  in  r' : e1 q' (s:ss)
  e1 w [] = e2 w
  e2 w = if w==0 then [] else let (w',y) = w `divMod` b in y : e2 w'
{-# LINE 2372 "asymm.lhs" #-}
decode10 :: [Int] -> [Symbol] 
decode10 = d0 0 where 
  d0 w (y:ys) | w < l      =  d0 (w * b + y) ys 
  d0 w ys                  =  d1 w ys 
  d1 w ys                  =  let  (q,r)  = w `divMod` total 
                                   s      = find r 
                                   w'     = count s * q + r - cumul s
                              in d2 s w' ys
  d2 s w (y:ys)  | w < l   =  d2 s (w * b + y) ys 
  d2 s w ys      | w >= l  = s : d1 w ys 
  d2 s w []                = []
{-# LINE 2400 "asymm.lhs" #-}
prop_EncodeDecode10 :: [CheckSymbol] -> Bool
prop_EncodeDecode10 xs = decode10 (encode10 t) == t
  where t = map unS xs
{-# LINE 2483 "asymm.lhs" #-}
runTests :: IO ()
runTests = sequence_ [
  quickCheck prop_EncodeDecode1,
  quickCheck prop_EncodeDecode2, 
  quickCheck prop_EncodeDecode3, 
  quickCheck prop_EncodeDecode4, 
  quickCheck prop_EconsumeDproduce5, quickCheck prop_EncodeDecode5, 
  quickCheck prop_EncodeDecode6,
  quickCheck prop_EncodeDecode7,
  quickCheck prop_EncodeDecode8,
  quickCheck prop_EncodeDecode9, 
  quickCheck prop_EncodeDecode10,
  quickCheck prop_Huffman
  ]