{-# LANGUAGE
    BangPatterns
  , UnboxedTuples
  , UnicodeSyntax
  #-}
-- | Lazy bitstrings based on 'L.ByteString' which treats a byte
-- stream as a bit sequence in the same way as WavPack's manner.
module Codec.Audio.WavPack.BitString
    ( -- * The BitString Type
      BitString

      -- * Construction
    , (∅)
    , empty
    , singleton

      -- * Basic Interface
    , cons
    , snoc
    , head
    , uncons
    , tail
    , null
    , length

      -- * Conversion from/to 'L.ByteString'
    , fromByteString
    , toByteString

      -- * Substrings
    , take
    , drop
    , splitAt
    )
    where
import qualified Data.ByteString.Lazy as L
import Data.Bits
import Data.Int
import qualified Data.Strict as S
import Data.Word
import Prelude hiding (drop, head, length, null, splitAt, tail, take)
import Prelude.Unicode

-- | The BitString Type.
data BitString
    = BitString {
        leftRem    ∷ !Remnant
      , leftBytes  ∷ !L.ByteString
      -- | reversed
      , rightBytes ∷ !L.ByteString
      , rightRem   ∷ !Remnant
      }
    deriving Show

-- | /O(n)/ The bitstrings must have finite lengths to test the
-- equality.
instance Eq BitString where
    a == b = leftRem    a' ≡ leftRem    b' ∧
             leftBytes  a' ≡ leftBytes  b' ∧
             rightRem   a' ≡ rightRem   b'
        where
          a' = normalise a
          b' = normalise b

normalise ∷ BitString → BitString
normalise bs
    | remLen (leftRem bs) ≡ 8
        = normalise $ bs {
                        leftRem   = remEmpty
                      , leftBytes = L.cons (remByte $ leftRem bs) $ leftBytes bs
                      }
    | remLen (rightRem bs) ≡ 8
        = normalise $ bs {
                        rightBytes = L.cons (remByte $ rightRem bs) $ rightBytes bs
                      , rightRem   = remEmpty
                      }
    | otherwise
        = bs {
            leftBytes  = leftBytes bs `L.append` (L.reverse $ rightBytes bs)
          , rightBytes = L.empty
          }

data Remnant
    = Remnant {
      -- | bit length from 0 to 8
        remLen  ∷ !Int
      -- | current byte @(0 <= remByte <= (1 `'shiftL'` 'remLen') - 1)@
      , remByte ∷ !Word8
      }
    deriving (Eq, Show)

remEmpty ∷ Remnant
remEmpty = Remnant 0 0

remNull ∷ Remnant → Bool
remNull = (0 ≡) ∘ remLen

remSingleton ∷ Bool -> Remnant
remSingleton b = Remnant {
                   remLen  = 1
                 , remByte = if b then 1 else 0
                 }

byteToRem ∷ Word8 → Remnant
byteToRem w = Remnant {
                remLen  = 8
              , remByte = w
              }

remToStr ∷ Remnant → L.ByteString
remToStr r
    | remNull r = L.empty
    | otherwise = L.singleton $ remByte r

consRem ∷ Bool → Remnant → (# Remnant, S.Maybe Word8 #)
consRem b r
    | remLen r ≡ 8 = (# remSingleton b, S.Just $ remByte r #)
    | remLen r ≡ 7 = (# remEmpty, S.Just w' #)
    | otherwise     = let !r' = r {
                                  remLen  = remLen r + 1
                                , remByte = w'
                                }
                      in
                        (# r', S.Nothing #)
    where
      w' = (remByte r `shiftL` 1)
           .|.
           if b then 1 else 0

snocRem ∷ Remnant → Bool → (# S.Maybe Word8, Remnant #)
snocRem r b
    | remLen r ≡ 8 = (# S.Just $ remByte r, remSingleton b #)
    | remLen r ≡ 7 = (# S.Just w', remEmpty #)
    | otherwise     = let !r' = r {
                                  remLen  = remLen r + 1
                                , remByte = w'
                                }
                      in
                        (# S.Nothing, r' #)
    where
      w' | b         = bit (remLen r) .|. remByte r
         | otherwise = remByte r

unconsRem ∷ Remnant → (# S.Maybe Bool, Remnant #)
unconsRem r
    | remNull r = (# S.Nothing, remEmpty #)
    | otherwise = let !b  = remByte r `testBit` 0
                      !r' = Remnant {
                              remLen  = remLen r - 1
                            , remByte = remByte r `shiftR` 1
                            }
                  in
                    (# S.Just b, r' #)

unconsBytes ∷ L.ByteString → (# S.Maybe Remnant, L.ByteString #)
unconsBytes bs
    = case L.uncons bs of
        Just (!w, bs')
            → (# S.Just (byteToRem w), bs' #)
        Nothing
            → (# S.Nothing, L.empty #)

appendRem ∷ Remnant → Remnant → Remnant
appendRem a b = Remnant {
                  remLen  = remLen a + remLen b
                , remByte = (remByte a `shiftL` remLen b) .|. remByte b
                }

splitRemAt ∷ Integral n ⇒ n → Remnant → (# S.Maybe Remnant, Remnant #)
splitRemAt n r
    | remLen r ≥ fromIntegral n
        = let !h  = Remnant {
                      remLen  = fromIntegral n
                    , remByte = remByte r .&. ((1 `shiftL` fromIntegral n) - 1)
                    }
              !r' = Remnant {
                      remLen  = remLen r - fromIntegral n
                    , remByte = remByte r `shiftR` fromIntegral n
                    }
          in
            (# S.Just h, r' #)
    | otherwise
        = (# S.Nothing, remEmpty #)

-- | /O(1)/ The same as 'empty'.
(∅) ∷ BitString
(∅) = empty

-- | /O(1)/ The empty 'BitString'.
empty ∷ BitString
empty = BitString {
          leftRem    = remEmpty
        , leftBytes  = L.empty
        , rightBytes = L.empty
        , rightRem   = remEmpty
        }

-- | /O(1)/ Convert a 'Bool' into a 'BitString'.
singleton ∷ Bool → BitString
singleton b = BitString {
                leftRem    = remSingleton b
              , leftBytes  = L.empty
              , rightBytes = L.empty
              , rightRem   = remEmpty
              }

-- | /O(1)/ Prepend a bit to the beginning of a 'BitString'.
cons ∷ Bool → BitString → BitString
cons b bs
    = case consRem b $ leftRem bs of
        (# lr', S.Just w #)
            → bs {
                 leftRem   = lr'
               , leftBytes = L.cons' w $ leftBytes bs
               }

        (# lr', S.Nothing #)
            → bs { leftRem = lr' }

-- | /O(1)/ Append a bit to the end of a 'BitString'.
snoc ∷ BitString → Bool → BitString
snoc bs b
    = case snocRem (rightRem bs) b of
        (# S.Just w, rr' #)
            → bs {
                 rightBytes = L.cons' w $ rightBytes bs
               , rightRem   = rr'
               }

        (# S.Nothing, rr' #)
            → bs { rightRem = rr' }

-- | /amortized O(1)/ Extract the first bit of a 'BitString', which
-- must be non-empty.
head ∷ BitString → Bool
head bs = case uncons bs of
            Just (b, _ )
                → b
            Nothing
                → error "head: empty BitString"

-- | /amortized O(1)/ Extract the bits after the head of a
-- 'BitString', which must be non-empty.
tail ∷ BitString → BitString
tail bs = case uncons bs of
            Just (_, bs')
                → bs'
            Nothing
                → error "tail: empty BitString"

-- | /amortized O(1)/ Extract the the head and tail of a 'BitString',
-- returning 'Nothing' if it's empty.
uncons ∷ BitString → Maybe (Bool, BitString)
uncons bs
    = case unconsRem $ leftRem bs of
        (# S.Just b, lr' #)
            → Just (b, bs { leftRem = lr' })

        (# S.Nothing, _ #)
            → case unconsBytes $ leftBytes bs of
                 (# S.Just lr, lb' #)
                     → let !bs' = bs { leftRem   = lr
                                      , leftBytes = lb'
                                      }
                        in
                          uncons bs'

                 (# S.Nothing, _ #)
                     → if L.null $ rightBytes bs then
                            case unconsRem $ rightRem bs of
                              (# S.Just b, rr' #)
                                  → Just (b, bs { rightRem = rr' })

                              (# S.Nothing, _ #)
                                  → Nothing
                        else
                            let !bs' = bs { leftBytes = L.reverse
                                                        $ rightBytes bs
                                          }
                            in
                              uncons bs'

-- | /O(1)/ Test whether a 'BitString' is empty.
null ∷ BitString → Bool
null bs
    = remNull (leftRem bs)
      ∧
      L.null (leftBytes bs)
      ∧
      L.null (rightBytes bs)
      ∧
      remNull (rightRem bs)

-- | /O(n)/ @'length' bs@ returns the number of bits in @bs@. @bs@
-- must have a finite length.
length ∷ Integral n ⇒ BitString → n
length bs
    = fromIntegral $ ( fromIntegral (remLen $ leftRem bs)
                       +
                       L.length (leftBytes bs) ⋅ 8
                       +
                       L.length (rightBytes bs) ⋅ 8
                       +
                       fromIntegral (remLen $ rightRem bs)
                     )
{-# SPECIALISE length ∷ BitString → Int64 #-}

-- | /O(1)/ Convert a 'L.ByteString' into a 'BitString'.
fromByteString ∷ L.ByteString → BitString
fromByteString bs = BitString {
                      leftRem    = remEmpty
                    , leftBytes  = bs
                    , rightBytes = L.empty
                    , rightRem   = remEmpty
                    }

-- | /O(n)/ @'toByteString' bs@ converts @bs@ into 'L.ByteString',
-- padding incomplete bytes to single bytes with necessary 0's at
-- their MSB. Thus the following equation does not hold when the
-- length of @bs@ isn't multiple of 8.
--
-- > fromByteString . toByteString = id
--
-- But the following always holds true.
--
-- > toByteString . fromByteString = id
--
-- Note that @bs@ must have a finite length.
toByteString ∷ BitString → L.ByteString
toByteString bs
    = L.concat [ remToStr $ leftRem bs
               , leftBytes bs
               , L.reverse $ rightBytes bs
               , remToStr $ rightRem bs
               ]

-- | /O(n)/ @'take' n bs@ returns the prefix of @bs@ of length
-- @n@. @bs@ must have at least @n@ bits.
take ∷ Integral n ⇒ n → BitString → BitString
take n xs = case splitAt n xs of
              (taken, _) → taken

-- | /O(n)/ @'drop' n bs@ returns the suffix of @bs@ after the first
-- @n@ bits. @bs@ must have at least @n@ bits.
drop ∷ Integral n ⇒ n → BitString → BitString
drop n xs = case splitAt n xs of
              (_, dropped) → dropped

-- | /O(n)/ @'splitAt' n bs@ is equivalent to @('take' n bs, 'drop' n
-- bs)@.
splitAt ∷ Integral n ⇒ n → BitString → (BitString, BitString)
splitAt n bs
    = case splitRemAt n $ leftRem bs of
        (# S.Just h, lr' #)
            -- The leftRem has enough bits.
            → let !h' = BitString {
                           leftRem    = h
                         , leftBytes  = L.empty
                         , rightBytes = L.empty
                         , rightRem   = remEmpty
                         }
               in
                 (h', bs { leftRem = lr' })

        (# S.Nothing, _ #)
            → let !bytesToTake  = (n - fromIntegral (remLen $ leftRem bs)) `div` 8
                  !amortised    = leftBytes bs `L.append` (L.reverse $ rightBytes bs)
                  -- 0 ≤ bitsToTake < 8 - remLen (leftRem bs)
                  -- ∴ remLen (leftRem bs) + bitsToTake < 8
                  !bitsToTake   = n - bytesToTake ⋅ 8 - fromIntegral (remLen $ leftRem bs)
              in
                case (# bytesToTake, L.uncons amortised #) of
                  (#_, Just (!w, amor') #)
                      -- There is at least one byte in the byte pool
                      -- (amortised).
                      → let !h1  = leftRem bs
                            (# h2, lr' #)
                                 = case splitRemAt bitsToTake $ byteToRem w of
                                    (# S.Just h2', lr'' #) → (# h2', lr'' #)
                                    (# S.Nothing , _    #) → error "internal error"
                            !h   = h1 `appendRem` h2
                            (bytes, lb')
                                 = L.splitAt (fromIntegral bytesToTake) amor'
                        in
                          if L.length bytes ≡ fromIntegral bytesToTake then
                              -- The byte pool actuall has at least
                              -- bytesToTake bytes.
                              let !h'  = BitString {
                                           leftRem    = h
                                         , leftBytes  = bytes
                                         , rightBytes = L.empty
                                         , rightRem   = remEmpty
                                         }
                                  !bs' = BitString {
                                           leftRem    = lr'
                                         , leftBytes  = lb'
                                         , rightBytes = L.empty
                                         , rightRem   = rightRem bs
                                         }
                              in
                                (h', bs')
                          else
                              error "splitAt: not enough bits"

                  (# 0, Nothing #)
                      -- No bytes are in the byte pool but the
                      -- rightRem may have enough bits.
                      → case splitRemAt bitsToTake $ rightRem bs of
                           (# S.Just h, rr' #)
                               → let !h'  = BitString {
                                              leftRem    = leftRem bs
                                            , leftBytes  = L.empty
                                            , rightBytes = L.empty
                                            , rightRem   = h
                                            }
                                     !bs' = BitString {
                                              leftRem    = remEmpty
                                            , leftBytes  = L.empty
                                            , rightBytes = L.empty
                                            , rightRem   = rr'
                                            }
                                 in
                                   (h', bs')

                           (# S.Nothing, _ #)
                               → error "splitAt: not enough bits"

                  (# _, Nothing #)
                      -- No bytes are in the byte pool but more than 8
                      -- bits are requested. The rightRem can't have
                      -- that many bits.
                      → error "splitAt: not enough bits"