{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE CPP #-}
module Codec.Picture.Jpg.Common
    ( DctCoefficients
    , JpgUnpackerParameter( .. )
    , decodeInt
    , dcCoefficientDecode
    , deQuantize
    , decodeRrrrSsss
    , zigZagReorderForward 
    , zigZagReorderForwardv
    , zigZagReorder
    , inverseDirectCosineTransform
    , unpackInt
    , unpackMacroBlock
    , rasterMap
    , decodeMacroBlock
    , decodeRestartInterval
    , toBlockSize
    ) where

#if !MIN_VERSION_base(4,8,0)
import Control.Applicative( pure, (<$>) )
#endif

import Control.Monad( when )
import Control.Monad.ST( ST, runST )
import Data.Bits( unsafeShiftL, unsafeShiftR, (.&.) )
import Data.Int( Int16, Int32 )
import Data.Maybe( fromMaybe )
import Data.Word( Word8 )
import qualified Data.Vector.Storable as VS
import qualified Data.Vector.Storable.Mutable as M
import Foreign.Storable ( Storable )

import Codec.Picture.Types
import Codec.Picture.BitWriter
import Codec.Picture.Jpg.Types
import Codec.Picture.Jpg.FastIdct
import Codec.Picture.Jpg.DefaultTable

-- | Same as for DcCoefficient, to provide nicer type signatures
type DctCoefficients = DcCoefficient

data JpgUnpackerParameter = JpgUnpackerParameter
    { dcHuffmanTree        :: !HuffmanPackedTree
    , acHuffmanTree        :: !HuffmanPackedTree
    , componentIndex       :: {-# UNPACK #-} !Int
    , restartInterval      :: {-# UNPACK #-} !Int
    , componentWidth       :: {-# UNPACK #-} !Int
    , componentHeight      :: {-# UNPACK #-} !Int
    , subSampling          :: !(Int, Int)
    , coefficientRange     :: !(Int, Int)
    , successiveApprox     :: !(Int, Int)
    , readerIndex          :: {-# UNPACK #-} !Int
      -- | When in progressive mode, we can have many
      -- color in a scan are only one. The indices changes
      -- on this fact, when mixed, there is whole 
      -- MCU for all color components, spanning multiple
      -- block lines. With only one color component we use
      -- the normal raster order.
    , indiceVector         :: {-# UNPACK #-} !Int
    , blockIndex           :: {-# UNPACK #-} !Int
    , blockMcuX            :: {-# UNPACK #-} !Int
    , blockMcuY            :: {-# UNPACK #-} !Int
    }
    deriving Show

toBlockSize :: Int -> Int
toBlockSize v = (v + 7) `div` 8

decodeRestartInterval :: BoolReader s Int32
decodeRestartInterval = return (-1) {-  do
  bits <- replicateM 8 getNextBitJpg
  if bits == replicate 8 True
     then do
         marker <- replicateM 8 getNextBitJpg
         return $ packInt marker
     else return (-1)
        -}

{-# INLINE decodeInt #-}
decodeInt :: Int -> BoolReader s Int32
decodeInt ssss = do
    signBit <- getNextBitJpg
    let dataRange = 1 `unsafeShiftL` fromIntegral (ssss - 1)
        leftBitCount = ssss - 1
    -- First following bits store the sign of the coefficient, and counted in
    -- SSSS, so the bit count for the int, is ssss - 1
    if signBit
       then (\w -> dataRange + fromIntegral w) <$> unpackInt leftBitCount
       else (\w -> 1 - dataRange * 2 + fromIntegral w) <$> unpackInt leftBitCount

decodeRrrrSsss :: HuffmanPackedTree -> BoolReader s (Int, Int)
decodeRrrrSsss tree = do
    rrrrssss <- huffmanPackedDecode tree
    let rrrr = (rrrrssss `unsafeShiftR` 4) .&. 0xF
        ssss =  rrrrssss .&. 0xF
    pure (fromIntegral rrrr, fromIntegral ssss)

dcCoefficientDecode :: HuffmanPackedTree -> BoolReader s DcCoefficient
dcCoefficientDecode dcTree = do
    ssss <- huffmanPackedDecode dcTree
    if ssss == 0
       then return 0
       else fromIntegral <$> decodeInt (fromIntegral ssss)

-- | Apply a quantization matrix to a macroblock
{-# INLINE deQuantize #-}
deQuantize :: MacroBlock Int16 -> MutableMacroBlock s Int16
           -> ST s (MutableMacroBlock s Int16)
deQuantize table block = update 0
    where update 64 = return block
          update i = do
              val <- block `M.unsafeRead` i
              let finalValue = val * (table `VS.unsafeIndex` i)
              (block `M.unsafeWrite` i) finalValue
              update $ i + 1

inverseDirectCosineTransform :: MutableMacroBlock s Int16
                             -> ST s (MutableMacroBlock s Int16)
inverseDirectCosineTransform mBlock =
    fastIdct mBlock >>= mutableLevelShift

zigZagOrder :: MacroBlock Int
zigZagOrder = makeMacroBlock $ concat
    [[ 0, 1, 5, 6,14,15,27,28]
    ,[ 2, 4, 7,13,16,26,29,42]
    ,[ 3, 8,12,17,25,30,41,43]
    ,[ 9,11,18,24,31,40,44,53]
    ,[10,19,23,32,39,45,52,54]
    ,[20,22,33,38,46,51,55,60]
    ,[21,34,37,47,50,56,59,61]
    ,[35,36,48,49,57,58,62,63]
    ]

zigZagReorderForwardv :: (Storable a, Num a) => VS.Vector a -> VS.Vector a
zigZagReorderForwardv vec = runST $ do
    v <- M.new 64
    mv <- VS.thaw vec
    zigZagReorderForward v mv >>= VS.freeze

zigZagOrderForward :: MacroBlock Int
zigZagOrderForward = VS.generate 64 inv
  where inv i = fromMaybe 0 $ VS.findIndex (i ==) zigZagOrder

zigZagReorderForward :: (Storable a)
                     => MutableMacroBlock s a
                     -> MutableMacroBlock s a
                     -> ST s (MutableMacroBlock s a)
{-# SPECIALIZE INLINE zigZagReorderForward :: MutableMacroBlock s Int32
                                           -> MutableMacroBlock s Int32
                                           -> ST s (MutableMacroBlock s Int32) #-}
{-# SPECIALIZE INLINE zigZagReorderForward :: MutableMacroBlock s Int16
                                           -> MutableMacroBlock s Int16
                                           -> ST s (MutableMacroBlock s Int16) #-}
{-# SPECIALIZE INLINE zigZagReorderForward :: MutableMacroBlock s Word8
                                           -> MutableMacroBlock s Word8
                                           -> ST s (MutableMacroBlock s Word8) #-}
zigZagReorderForward zigzaged block = ordering zigZagOrderForward >> return zigzaged
  where ordering !table = reorder (0 :: Int)
          where reorder !i | i >= 64 = return ()
                reorder i  = do
                     let idx = table `VS.unsafeIndex` i
                     v <- block `M.unsafeRead` idx
                     (zigzaged `M.unsafeWrite` i) v
                     reorder (i + 1)

zigZagReorder :: MutableMacroBlock s Int16 -> MutableMacroBlock s Int16
              -> ST s (MutableMacroBlock s Int16)
zigZagReorder zigzaged block = do
    let update i =  do
            let idx = zigZagOrder `VS.unsafeIndex` i
            v <- block `M.unsafeRead` idx
            (zigzaged `M.unsafeWrite` i) v

        reorder 63 = update 63
        reorder i  = update i >> reorder (i + 1)

    reorder (0 :: Int)
    return zigzaged

-- | Unpack an int of the given size encoded from MSB to LSB.
unpackInt :: Int -> BoolReader s Int32
unpackInt = getNextIntJpg

{-# INLINE rasterMap #-}
rasterMap :: (Monad m)
          => Int -> Int -> (Int -> Int -> m ())
          -> m ()
rasterMap width height f = liner 0
  where liner y | y >= height = return ()
        liner y = columner 0
          where columner x | x >= width = liner (y + 1)
                columner x = f x y >> columner (x + 1)

pixelClamp :: Int16 -> Word8
pixelClamp n = fromIntegral . min 255 $ max 0 n

-- | Given a size coefficient (how much a pixel span horizontally
-- and vertically), the position of the macroblock, return a list
-- of indices and value to be stored in an array (like the final
-- image)
unpackMacroBlock :: Int    -- ^ Component count
                 -> Int -- ^ Width coefficient
                 -> Int -- ^ Height coefficient
                 -> Int -- ^ Component index
                 -> Int -- ^ x
                 -> Int -- ^ y
                 -> MutableImage s PixelYCbCr8
                 -> MutableMacroBlock s Int16
                 -> ST s ()
unpackMacroBlock compCount wCoeff hCoeff compIdx x y
                 (MutableImage { mutableImageWidth = imgWidth,
                                 mutableImageHeight = imgHeight, mutableImageData = img })
                 block = rasterMap dctBlockSize dctBlockSize unpacker
  where unpacker i j = do
          let yBase = y * dctBlockSize + j * hCoeff
          compVal <- pixelClamp <$> (block `M.unsafeRead` (i + j * dctBlockSize))
          rasterMap wCoeff hCoeff $ \wDup hDup -> do
             let xBase = x * dctBlockSize + i * wCoeff
                 xPos = xBase + wDup
                 yPos = yBase + hDup

             when (xPos < imgWidth && yPos < imgHeight)
                  (do let mutableIdx = (xPos + yPos * imgWidth) * compCount + compIdx
                      (img `M.unsafeWrite` mutableIdx) compVal)

-- | This is one of the most important function of the decoding,
-- it form the barebone decoding pipeline for macroblock. It's all
-- there is to know for macro block transformation
decodeMacroBlock :: MacroBlock DctCoefficients
                 -> MutableMacroBlock s Int16
                 -> MutableMacroBlock s Int16
                 -> ST s (MutableMacroBlock s Int16)
decodeMacroBlock quantizationTable zigZagBlock block =
    deQuantize quantizationTable block >>= zigZagReorder zigZagBlock
                                       >>= inverseDirectCosineTransform