{-# LANGUAGE
    BangPatterns
  , FlexibleContexts
  , ScopedTypeVariables
  , UnicodeSyntax
  #-}
{-| This module provides entropy word encoding and decoding functions
using a variation on the Rice method.  This was introduced in wavpack
3.93 because it allows splitting the data into a \"lossy\" stream and
a \"correction\" stream in a very efficient manner and is therefore
ideal for the "hybrid" mode.  For 4.0, the efficiency of this method
was significantly improved by moving away from the normal Rice
restriction of using powers of two for the modulus divisions and now
the method can be used for both hybrid and pure lossless encoding.

Samples are divided by median probabilities at 5\/7 (71.43%), 10\/49
(20.41%), and 20\/343 (5.83%). Each zone has 3.5 times fewer samples
than the previous. Using standard Rice coding on this data would
result in 1.4 bits per sample average (not counting sign
bit). However, there is a very simple encoding that is over 99%
efficient with this data and results in about 1.22 bits per sample. -}
module Codec.Audio.WavPack.Words
    ( WordsData(..)

    , getWordsLossless
    )
    where
import Codec.Audio.WavPack.Entropy
import Codec.Audio.WavPack.Internal
import Control.Monad.Cont
import Control.Monad.ST
import Control.Monad.Trans
import Control.Monad.Unicode
import Data.Bits
import Data.Bitstream.Generic (Bitstream)
import qualified Data.Bitstream.Generic as B
import Data.Int
import Data.STRef
import qualified Data.Vector.Generic.Mutable as MV
import Data.Word
import Prelude hiding (break)
import Prelude.Unicode

-- | FIXME
data WordsData s
    = WordsData {
        wdBitrateDelta ∷ !(STRef s (Word32, Word32))
      , wdBitrateAcc   ∷ !(STRef s (Word32, Word32))
      , wdPendingData  ∷ !(STRef s Word32)
      , wdHoldingOne   ∷ !(STRef s Word32)
      , wdZeroesAcc    ∷ !(STRef s Word32)
      , wdHoldingZero  ∷ !(STRef s Bool)
      , wdPendingCount ∷ !(STRef s Int)
      , wdEntropyData  ∷ !(EntropyData s, EntropyData s)
      }

-- | This is an optimized version of 'getWord' that is used for
-- lossless only ('edErrorLimit' ≡ 0). Also, rather than obtaining a
-- single sample, it can be used to obtain an entire buffer of either
-- mono or stereo samples.
getWordsLossless ∷ ∀bs v s. (Bitstream bs, MV.MVector v Int32)
                 ⇒ Bool       -- ^ Is the stream monaural?
                 → WordsData s
                 → STRef s bs -- ^ WV bitstream
                 → Int        -- ^ Number of samples to get
                 → ST s (v s Int32)
{-# INLINEABLE getWordsLossless #-}
getWordsLossless isMono w bs nSamples0
    = do v ← MV.new nSamples
         n ← runContT (for 0 (< nSamples) (+ 1) (loop v)) return
         return $ MV.take n v
    where
      nSamples ∷ Int
      nSamples = if isMono
                 then nSamples0
                 else nSamples0 ⋅ 2

      loop ∷ v s Int32
           → Int
           → ContT Int (ST s) ()
           → ContT Int (ST s) ()
           → ContT Int (ST s) ()
      loop v n break continue
          = do let c | isMono        = fst $ wdEntropyData w
                     | n `testBit` 0 = fst $ wdEntropyData w
                     | otherwise     = snd $ wdEntropyData w
               med00   ← lift $ readSTRef (edMedian0 $ fst $ wdEntropyData w)
               hldZero ← lift $ readSTRef (wdHoldingZero w)
               hldOne  ← lift $ readSTRef (wdHoldingOne  w)
               med10   ← lift $ readSTRef (edMedian0 $ snd $ wdEntropyData w)
               when (med00 < 2 ∧ hldZero ≡ False ∧ hldOne ≡ 0 ∧ med10 < 2) $
                    do zAcc ← lift $ readSTRef (wdZeroesAcc w)
                       if zAcc > 0 then
                           do lift $ modifySTRef (wdZeroesAcc w) ((-) 1)
                              when (zAcc > 1) $
                                   do lift $ MV.unsafeWrite v n 0
                                      continue
                         else
                           do cBits ← lift $ takeWhileLessThan id 33 bs

                              when (cBits ≡ 33) $
                                   break

                              if cBits < 2 then
                                  lift $ writeSTRef (wdZeroesAcc w) cBits
                                else
                                  error "FIXME"
               error "FIXME"

{-
getWordsLossless ∷ ∀bs v. (Bitstream bs, GV.Vector v Int32)
                 ⇒ Bool -- ^ Is the stream monaural?
                 → WordsData
                 → bs   -- ^ WV bitstream.
                 → Int   -- ^ Number of samples to get.
                 → (# WordsData, bs, v Int32 #)
{-# INLINEABLE getWordsLossless #-}
getWordsLossless isMono w0 bs0 nSamples0
    = let v0  = New.create $ MV.new nSamples
          (# w1, bs1, n1, v1 #)
              = go0 w0 bs0 0 v0
          v2  = GV.new $ New.take n1 v1
      in
        (# w1, bs1, v2 #)
    where
      go0 ∷ WordsData → bs → Int → New v Int32
          → (# WordsData, bs, Int, New v Int32 #)
      go0 w bs n v
          | n ≥ nSamples
              = (# w, bs, n, v #)
          | edMedian0 (fst $ wdEntropyData w) < 2 ∧
            wdHoldingZero w ≡ False               ∧
            wdHoldingOne  w ≡ 0                   ∧
            edMedian1 (fst $ wdEntropyData w) < 2
              = if wdZeroesAcc w > 0 then
                    let w' = w { wdZeroesAcc = wdZeroesAcc w - 1 }
                    in
                      if wdZeroesAcc w' > 0 then
                          let (# n', v' #) = appendWord 0 n v
                          in
                            go0 w' bs n' v'
                      else
                          go1 w' bs n v
                else
                    let cBits = min 33 $ B.length (B.takeWhile id bs)
                        bs'   = B.drop cBits bs
                    in
                      if cBits ≡ 33 then
                          (# w, bs', n, v #)
                      else
                          let (# w', bs'' #) = go0' cBits w bs'
                          in
                            if wdZeroesAcc w' > 0 then
                                let w'' = w' {
                                            wdEntropyData =
                                                ( clearMedian $ fst $ wdEntropyData w'
                                                , clearMedian $ snd $ wdEntropyData w' )
                                          }
                                    (# n', v' #)
                                        = appendWord 0 n v
                                in
                                  go0 w'' bs'' n' v'
                            else
                                go1 w' bs'' n v
          | otherwise
              = go1 w bs n v

      go0' ∷ Word32 → WordsData → bs → (# WordsData, bs #)
      go0' cBits w bs
          | cBits < 2
              = let w' = w { wdZeroesAcc = cBits }
                in
                  (# w', bs #)
          | otherwise
              = let w' = w { wdZeroesAcc = 0 }
                in
                  go0'' 1 cBits w' bs

      go0'' ∷ Word32 → Word32 → WordsData → bs → (# WordsData, bs #)
      go0'' mask cBits w bs
          | cBits ≡ 1
              = let w' = w { wdZeroesAcc = wdZeroesAcc w .|. mask }
                in
                  (# w', bs #)
          | otherwise
              = let cBits' = cBits - 1
                    w'     = if B.head bs then
                                 w { wdZeroesAcc = wdZeroesAcc w .|. mask }
                             else
                                 w
                    mask'  = mask `shiftL` 1
                    bs'    = B.tail bs
                in
                  go0'' mask' cBits' w' bs'

      go1 ∷ WordsData → bs → Int → New v Int32
          → (# WordsData, bs, Int, New v Int32 #)
      go1 w bs n v
          | wdHoldingZero w
              = let w' = w { wdHoldingZero = False }
                in
                  go2 0 w' bs n v
          | otherwise
              = let next8 ∷ Word8
                    next8 = B.toBits (B.take (8 ∷ Int) bs)
                in
                  if next8 ≡ 0xFF then
                      error "FIXME"
                  else
                      error "FIXME"

      go2 ∷ Word32 → WordsData → bs → Int → New v Int32
          → (# WordsData, bs, Int, New v Int32 #)
      go2 0 w bs n v
          = let ent  = getEntropy n w
                low  = 0
                high = getMedian0 ent - 1
                ent' = decMedian0 ent
                w'   = setEntropy ent' n w
            in
              go3 low high w' bs n v
      go2 1 w bs n v
          = let ent  = getEntropy n w
                low  = getMedian0 ent
                high = low + getMedian1 ent - 1
                ent' = (incMedian0 ∘ decMedian1) ent
                w'   = setEntropy ent' n w
            in
              go3 low high w' bs n v
      go2 2 w bs n v
          = let ent   = getEntropy n w
                low   = getMedian0 ent + getMedian1 ent
                high  = low + getMedian2 ent - 1
                ent'  = (incMedian0 ∘ incMedian1 ∘ decMedian2) ent
                w'    = setEntropy ent' n w
            in
              go3 low high w' bs n v
      go2 onesCount w bs n v
          = let ent   = getEntropy n w
                low   = getMedian0 ent + getMedian1 ent + (onesCount-2) ⋅ getMedian2 ent
                high  = low + getMedian2 ent - 1
                ent'  = (incMedian0 ∘ incMedian1 ∘ incMedian2) ent
                w'    = setEntropy ent' n w
            in
              go3 low high w' bs n v

      go3 ∷ Word32 → Word32 → WordsData → bs → Int → New v Int32
          → (# WordsData, bs, Int, New v Int32 #)
      go3 low high w bs n v
          = let (# code, bs' #)
                     = readCode bs (high - low)
                low' = low + code
                word = if B.head bs' then
                           fromIntegral $ complement low'
                       else
                           fromIntegral low'
                bs'' = B.tail bs'
                (# n', v' #)
                     = appendWord word n v
            in
              go0 w bs'' n' v'

      appendWord ∷ Int32 → Int → New v Int32 → (# Int, New v Int32 #)
      appendWord word n v
          = let v' = New.modify (\mv → MV.unsafeWrite mv n word) v
                n' = n + 1
            in
              (# n', v' #)

      getEntropy ∷ Int → WordsData → EntropyData
      getEntropy n w
          | isMono        = fst $ wdEntropyData w
          | n `testBit` 0 = fst $ wdEntropyData w
          | otherwise     = snd $ wdEntropyData w

      setEntropy ∷ EntropyData → Int → WordsData → WordsData
      setEntropy e n w
          | isMono        = w { wdEntropyData = (e, snd $ wdEntropyData w) }
          | n `testBit` 0 = w { wdEntropyData = (e, snd $ wdEntropyData w) }
          | otherwise     = w { wdEntropyData = (fst $ wdEntropyData w, e) }
-}

-- | Read a single unsigned value from the specified bitstream with a
-- value from 0 to maxCode. If there are exactly a power of two number
-- of possible codes then this will read a fixed number of bits;
-- otherwise it reads the minimum number of bits and then determines
-- whether another bit is needed to define the code.
readCode ∷ Bitstream bs ⇒ STRef s bs → Word32 → ST s Word32
{-# INLINEABLE readCode #-}
readCode bs 0       = return 0
readCode bs 1       = fmap b2n $ takeHead bs
readCode bs maxCode
    = do let bitCount = countBits maxCode
             extras   = bit bitCount - maxCode - 1
         code ← takeBits (bitCount - 1) bs
         if code ≥ extras then
             do nextBit ← takeHead bs
                return $ (code `shiftL` 1) - extras + b2n nextBit
           else
             return code

takeHead ∷ Bitstream bs ⇒ STRef s bs → ST s Bool
{-# INLINEABLE takeHead #-}
takeHead bsr
    = do bs ← readSTRef bsr
         writeSTRef bsr (B.tail bs)
         return (B.head bs)

takeWhileLessThan ∷ (Integral n, Bitstream bs)
                  ⇒ (Bool → Bool)
                  → n
                  → STRef s bs
                  → ST s n
{-# INLINEABLE takeWhileLessThan #-}
takeWhileLessThan f n bsr = go 0
    where
      {-# INLINE go #-}
      go i | i < n
               = do b ← takeHead bsr
                    if b then
                        go (i + 1)
                      else
                        return i
           | otherwise
               = return i

takeBits ∷ (Integral n, Bitstream bs, Bits a) ⇒ n → STRef s bs → ST s a
{-# INLINEABLE takeBits #-}
takeBits n bsr
    = do bs ← readSTRef bsr
         writeSTRef bsr (B.drop n bs)
         return (B.toBits (B.take n bs))

-- | C style /for/ loop with /break/ and /continue/.
for ∷ ∀m α. MonadCont m
    ⇒ α          -- ^ Initial state
    → (α → Bool) -- ^ Continue-the-loop predicate
    → (α → α)    -- ^ State modifier
    → (α → m () → m () → m ()) -- ^ Loop body taking breaker and
                               -- continuer
    → m α        -- ^ Final state
for α0 contLoop next body
    = callCC $ \break → loop break α0
    where
      loop ∷ (α → m ()) → α → m α
      loop break α
          | contLoop α
              = do callCC $ \continue → body α (break α) (continue ())
                   loop break (next α)
          | otherwise
              = return α