content: {-# LANGUAGE BangPatterns, CPP #-}
  module Codec.Picture.Gif.Internal.LZWEncoding( lzwEncode ) where
  
  #if !MIN_VERSION_base(4,8,0)
  import Control.Applicative( (<$>) )
  import Data.Monoid( mempty )
  #endif
  
  import Control.Monad.ST( runST )
  import qualified Data.ByteString.Lazy as L
  import Data.Maybe( fromMaybe )
  import Data.Word( Word8 )
  
  #if MIN_VERSION_containers(0,5,0)
  import qualified Data.IntMap.Strict as I
  #else
  import qualified Data.IntMap as I
  #endif
  import qualified Data.Vector.Storable as V
  
  import Codec.Picture.BitWriter
  
  type Trie = I.IntMap TrieNode
  
  data TrieNode = TrieNode
      { trieIndex :: {-# UNPACK #-} !Int
      , trieSub   :: !Trie
      }
  
  emptyNode :: TrieNode
  emptyNode = TrieNode
      { trieIndex = -1
      , trieSub = mempty
      }
  
  initialTrie :: Trie
  initialTrie = I.fromList
      [(i, emptyNode { trieIndex = i }) | i <- [0 .. 255]]
  
  lookupUpdate :: V.Vector Word8 -> Int -> Int -> Trie -> (Int, Int, Trie)
  lookupUpdate vector freeIndex firstIndex trie =
      matchUpdate $ go trie 0 firstIndex 
    where
      matchUpdate (lzwOutputIndex, nextReadIndex, sub) =
          (lzwOutputIndex, nextReadIndex, fromMaybe trie sub)
  
      maxi = V.length vector
      go !currentTrie !prevIndex !index
        | index >= maxi = (prevIndex, index, Nothing)
        | otherwise = case I.lookup val currentTrie of
            Just (TrieNode ix subTable) ->
                let (lzwOutputIndex, nextReadIndex, newTable) =
                          go subTable ix $ index + 1
                    tableUpdater t =
                        I.insert val (TrieNode ix t) currentTrie
                in
                (lzwOutputIndex, nextReadIndex, tableUpdater <$> newTable)
  
            Nothing | index == maxi -> (prevIndex, index, Nothing)
                    | otherwise -> (prevIndex, index, Just $ I.insert val newNode currentTrie)
  
        where val = fromIntegral $ vector `V.unsafeIndex` index
              newNode = emptyNode { trieIndex = freeIndex }
  
  lzwEncode :: Int -> V.Vector Word8 -> L.ByteString
  lzwEncode initialKeySize vec = runST $ do
      bitWriter <- newWriteStateRef 
  
      let updateCodeSize 12 writeIdx _
              | writeIdx == 2 ^ (12 :: Int) - 1 = do
                 writeBitsGif bitWriter (fromIntegral clearCode) 12
                 return (startCodeSize, firstFreeIndex, initialTrie)
  
          updateCodeSize codeSize writeIdx trie
              | writeIdx == 2 ^ codeSize =
                  return (codeSize + 1, writeIdx + 1, trie)
              | otherwise = return (codeSize, writeIdx + 1, trie)
  
          go readIndex (codeSize, _, _) | readIndex >= maxi =
              writeBitsGif bitWriter (fromIntegral endOfInfo) codeSize
          go !readIndex (!codeSize, !writeIndex, !trie) = do
              let (indexToWrite, endIndex, trie') =
                      lookuper writeIndex readIndex trie
              writeBitsGif bitWriter (fromIntegral indexToWrite) codeSize
              updateCodeSize codeSize writeIndex trie'
                  >>= go endIndex 
  
      writeBitsGif bitWriter (fromIntegral clearCode) startCodeSize
      go 0 (startCodeSize, firstFreeIndex, initialTrie)
  
      finalizeBoolWriterGif bitWriter
    where
      maxi = V.length vec
  
      startCodeSize = initialKeySize + 1
  
      clearCode = 2 ^ initialKeySize :: Int
      endOfInfo = clearCode + 1
      firstFreeIndex = endOfInfo + 1
      
      lookuper = lookupUpdate vec
  