content: {-# LANGUAGE CPP #-}
  module Codec.Picture.Gif.Internal.LZW( decodeLzw, decodeLzwTiff ) where
  
  #if !MIN_VERSION_base(4,8,0)
  import Control.Applicative( (<$>) )
  #endif
  
  import Data.Word( Word8 )
  import Control.Monad( when, unless )
  
  import Data.Bits( (.&.) )
  
  import Control.Monad.ST( ST )
  import Control.Monad.Trans.Class( MonadTrans, lift )
  
  import Foreign.Storable ( Storable )
  
  import qualified Data.ByteString as B
  import qualified Data.Vector.Storable.Mutable as M
  
  import Codec.Picture.BitWriter
  
  {-# INLINE (.!!!.) #-}
  (.!!!.) :: (Storable a) => M.STVector s a -> Int -> ST s a
  (.!!!.) = M.unsafeRead 
          {-M.read-}
  
  {-# INLINE (..!!!..) #-}
  (..!!!..) :: (MonadTrans t, Storable a)
            => M.STVector s a -> Int -> t (ST s) a
  (..!!!..) v idx = lift $ v .!!!. idx
  
  {-# INLINE (.<-.) #-}
  (.<-.) :: (Storable a) => M.STVector s a -> Int -> a -> ST s ()
  (.<-.) = M.unsafeWrite 
           {-M.write-}
  
  {-# INLINE (..<-..) #-}
  (..<-..) :: (MonadTrans t, Storable a)
           => M.STVector s a -> Int -> a -> t (ST s) ()
  (..<-..) v idx = lift . (v .<-. idx)
  
  
  duplicateData :: (MonadTrans t, Storable a)
                => M.STVector s a -> M.STVector s a
                -> Int -> Int -> Int -> t (ST s) ()
  duplicateData src dest sourceIndex size destIndex = lift $ aux sourceIndex destIndex
    where endIndex = sourceIndex + size
          aux i _ | i == endIndex  = return ()
          aux i j = do
            src .!!!. i >>= (dest .<-. j)
            aux (i + 1) (j + 1)
  
  rangeSetter :: (Storable a, Num a)
              => Int -> M.STVector s a
              -> ST s (M.STVector s a)
  rangeSetter count vec = aux 0
    where aux n | n == count = return vec
          aux n = (vec .<-. n) (fromIntegral n) >> aux (n + 1)
  
  decodeLzw :: B.ByteString -> Int -> Int -> M.STVector s Word8
            -> BoolReader s ()
  decodeLzw str maxBitKey initialKey outVec = do
      setDecodedString str
      lzw GifVariant maxBitKey initialKey 0 outVec
  
  isOldTiffLZW :: B.ByteString -> Bool
  isOldTiffLZW str = firstByte == 0 && secondByte == 1
      where firstByte = str `B.index` 0
            secondByte = (str `B.index` 1) .&. 1
  
  decodeLzwTiff :: B.ByteString -> M.STVector s Word8 -> Int
                -> BoolReader s()
  decodeLzwTiff str outVec initialWriteIdx = do
      if isOldTiffLZW str then
        setDecodedString str
      else
        setDecodedStringMSB str
      let variant | isOldTiffLZW str = OldTiffVariant
                  | otherwise = TiffVariant
      lzw variant 12 9 initialWriteIdx outVec
  
  data TiffVariant =
        GifVariant
      | TiffVariant
      | OldTiffVariant
      deriving Eq
  
  -- | Gif image constraint from spec-gif89a, code size max : 12 bits.
  lzw :: TiffVariant -> Int -> Int -> Int -> M.STVector s Word8
      -> BoolReader s ()
  lzw variant nMaxBitKeySize initialKeySize initialWriteIdx outVec = do
      -- Allocate buffer of maximum size.
      lzwData <- lift (M.replicate maxDataSize 0) >>= resetArray
      lzwOffsetTable <- lift (M.replicate tableEntryCount 0) >>= resetArray
      lzwSizeTable <- lift $ M.replicate tableEntryCount 0
      lift $ lzwSizeTable `M.set` 1
  
      let firstVal code = do
              dataOffset <- lzwOffsetTable ..!!!.. code
              lzwData ..!!!.. dataOffset
  
          writeString at code = do
              dataOffset <- lzwOffsetTable ..!!!.. code
              dataSize   <- lzwSizeTable   ..!!!.. code
  
              when (at + dataSize <= maxWrite) $
                   duplicateData lzwData outVec dataOffset dataSize at
  
              return dataSize
  
          addString pos at code val = do
              dataOffset <- lzwOffsetTable ..!!!.. code
              dataSize   <- lzwSizeTable   ..!!!.. code
  
              when (pos < tableEntryCount) $ do
                (lzwOffsetTable ..<-.. pos) at
                (lzwSizeTable ..<-.. pos) $ dataSize + 1
  
              when (at + dataSize + 1 <= maxDataSize) $ do
                duplicateData lzwData lzwData dataOffset dataSize at
                (lzwData ..<-.. (at + dataSize)) val
  
              return $ dataSize + 1
  
          maxWrite = M.length outVec
          loop outWriteIdx writeIdx dicWriteIdx codeSize oldCode code
            | outWriteIdx >= maxWrite = return ()
            | code == endOfInfo = return ()
            | code == clearCode = do
                toOutput <- getNextCode startCodeSize
                unless (toOutput == endOfInfo) $ do
                  dataSize <- writeString outWriteIdx toOutput
                  getNextCode startCodeSize >>=
                    loop (outWriteIdx + dataSize)
                         firstFreeIndex firstFreeIndex startCodeSize toOutput
  
            | otherwise =  do
                (written, dicAdd) <-
                     if code >= writeIdx then do
                       c <- firstVal oldCode
                       wroteSize <- writeString outWriteIdx oldCode
                       (outVec ..<-.. (outWriteIdx + wroteSize)) c
                       addedSize <- addString writeIdx dicWriteIdx oldCode c
                       return (wroteSize + 1, addedSize)
                     else do
                       wroteSize <- writeString outWriteIdx code
                       c <- firstVal code
                       addedSize <- addString writeIdx dicWriteIdx oldCode c
                       return (wroteSize, addedSize)
  
                let new_code_size = updateCodeSize codeSize $ writeIdx + 1
                getNextCode new_code_size >>=
                  loop (outWriteIdx + written)
                       (writeIdx + 1)
                       (dicWriteIdx + dicAdd)
                       new_code_size
                       code
  
      getNextCode startCodeSize >>=
          loop initialWriteIdx firstFreeIndex firstFreeIndex startCodeSize 0
  
    where tableEntryCount =  2 ^ min 12 nMaxBitKeySize
          maxDataSize = tableEntryCount `div` 2 * (1 + tableEntryCount) + 1
  
          isNewTiff = variant == TiffVariant
          (switchOffset,  isTiffVariant) = case variant of
              GifVariant -> (0, False)
              TiffVariant -> (1, True)
              OldTiffVariant -> (0, True)
  
          initialElementCount = 2 ^ initialKeySize :: Int
          clearCode | isTiffVariant = 256
                    | otherwise = initialElementCount
  
          endOfInfo | isTiffVariant = 257
                    | otherwise = clearCode + 1
  
          startCodeSize 
                    | isTiffVariant = initialKeySize
                    | otherwise = initialKeySize + 1
  
          firstFreeIndex = endOfInfo + 1
  
          resetArray a = lift $ rangeSetter initialElementCount a
  
          updateCodeSize codeSize writeIdx
              | writeIdx == 2 ^ codeSize - switchOffset = min 12 $ codeSize + 1
              | otherwise = codeSize
  
          getNextCode s 
              | isNewTiff = fromIntegral <$> getNextBitsMSBFirst s
              | otherwise = fromIntegral <$> getNextBitsLSBFirst s
  
  