content: {-# LANGUAGE DeriveGeneric #-}
  {-# LANGUAGE TupleSections #-}
  {-# LANGUAGE FlexibleContexts #-}
  -- | Module used by the jpeg decoder internally, shouldn't be used
  -- in user code.
  module Codec.Picture.Jpg.Internal.DefaultTable( DctComponent( .. )
                                       , HuffmanTree( .. )
                                       , HuffmanTable
                                       , HuffmanPackedTree
                                       , MacroBlock
                                       , QuantificationTable
                                       , HuffmanWriterCode 
                                       , scaleQuantisationMatrix
                                       , makeMacroBlock
                                       , makeInverseTable
                                       , buildHuffmanTree
                                       , packHuffmanTree
                                       , huffmanPackedDecode
  
                                       , defaultChromaQuantizationTable
  
                                       , defaultLumaQuantizationTable
  
                                       , defaultAcChromaHuffmanTree
                                       , defaultAcChromaHuffmanTable
  
                                       , defaultAcLumaHuffmanTree 
                                       , defaultAcLumaHuffmanTable 
  
                                       , defaultDcChromaHuffmanTree 
                                       , defaultDcChromaHuffmanTable
  
                                       , defaultDcLumaHuffmanTree
                                       , defaultDcLumaHuffmanTable
                                       ) where
  
  import Control.DeepSeq( NFData(..) )
  import Data.Int( Int16 )
  import Foreign.Storable ( Storable )
  import Control.Monad.ST( runST )
  import qualified Data.Vector.Storable as SV
  import qualified Data.Vector as V
  import Data.Bits( unsafeShiftL, (.|.), (.&.) )
  import Data.Word( Word8, Word16 )
  import Data.List( foldl' )
  import qualified Data.Vector.Storable.Mutable as M
  import GHC.Generics( Generic )
  
  import Codec.Picture.BitWriter
  
  -- | Tree storing the code used for huffman encoding.
  data HuffmanTree = Branch HuffmanTree HuffmanTree -- ^ If bit is 0 take the first subtree, if 1, the right.
                   | Leaf Word8       -- ^ We should output the value
                   | Empty            -- ^ no value present
                   deriving (Eq, Show)
  
  type HuffmanPackedTree = SV.Vector Word16
  
  type HuffmanWriterCode = V.Vector (Word8, Word16)
  
  packHuffmanTree :: HuffmanTree -> HuffmanPackedTree
  packHuffmanTree tree = runST $ do
      table <- M.replicate 512 0x8000
      let aux (Empty) idx = return $ idx + 1
          aux (Leaf v) idx = do
              (table `M.unsafeWrite` idx) $ fromIntegral v .|. 0x4000
              return $ idx + 1
  
          aux (Branch i1@(Leaf _) i2@(Leaf _)) idx =
              aux i1 idx >>= aux i2
  
          aux (Branch i1@(Leaf _) i2) idx = do
              _ <- aux i1 idx
              ix2 <- aux i2 $ idx + 2
              (table `M.unsafeWrite` (idx + 1)) $ fromIntegral $ idx + 2
              return ix2
  
          aux (Branch i1 i2@(Leaf _)) idx = do
              ix1 <- aux i1 (idx + 2)
              _ <- aux i2 (idx + 1)
              (table `M.unsafeWrite` idx) . fromIntegral $ idx + 2
              return ix1
  
          aux (Branch i1 i2) idx = do
              ix1 <- aux i1 (idx + 2)
              ix2 <- aux i2 ix1
              (table `M.unsafeWrite` idx) (fromIntegral $ idx + 2)
              (table `M.unsafeWrite` (idx + 1)) (fromIntegral ix1)
              return ix2
      _ <- aux tree 0
      SV.unsafeFreeze table
  
  makeInverseTable :: HuffmanTree -> HuffmanWriterCode
  makeInverseTable t = V.replicate 255 (0,0) V.// inner 0 0 t
    where inner _     _     Empty   = []
          inner depth code (Leaf v) = [(fromIntegral v, (depth, code))]
          inner depth code (Branch l r) =
            inner (depth + 1) shifted l ++ inner (depth + 1) (shifted .|. 1) r
              where shifted = code `unsafeShiftL` 1
  
  -- | Represent a compact array of 8 * 8 values. The size
  -- is not guarenteed by type system, but if makeMacroBlock is
  -- used, everything should be fine size-wise
  type MacroBlock a = SV.Vector a
  
  type QuantificationTable = MacroBlock Int16
  
  -- | Helper function to create pure macro block of the good size.
  makeMacroBlock :: (Storable a) => [a] -> MacroBlock a
  makeMacroBlock = SV.fromListN 64
  
  -- | Enumeration used to search in the tables for different components.
  data DctComponent = DcComponent | AcComponent
      deriving (Eq, Show, Generic)
  instance NFData DctComponent
  
  -- | Transform parsed coefficients from the jpeg header to a
  -- tree which can be used to decode data.
  buildHuffmanTree :: [[Word8]] -> HuffmanTree
  buildHuffmanTree table = foldl' insertHuffmanVal Empty
                         . concatMap (\(i, t) -> map (i + 1,) t)
                         $ zip ([0..] :: [Int]) table
    where isTreeFullyDefined Empty = False
          isTreeFullyDefined (Leaf _) = True
          isTreeFullyDefined (Branch l r) = isTreeFullyDefined l && isTreeFullyDefined r
  
          insertHuffmanVal Empty (0, val) = Leaf val
          insertHuffmanVal Empty (d, val) = Branch (insertHuffmanVal Empty (d - 1, val)) Empty
          insertHuffmanVal (Branch l r) (d, val)
              | isTreeFullyDefined l = Branch l (insertHuffmanVal r (d - 1, val))
              | otherwise            = Branch (insertHuffmanVal l (d - 1, val)) r
          insertHuffmanVal (Leaf _) _ = error "Inserting in value, shouldn't happen"
  
  scaleQuantisationMatrix :: Int -> QuantificationTable -> QuantificationTable 
  scaleQuantisationMatrix quality
      | quality < 0 = scaleQuantisationMatrix 0
          -- shouldn't show much difference than with 1,
          -- but hey, at least we're complete
      | quality == 0 = SV.map (scale (10000 :: Int))
      | quality < 50 = let qq = 5000 `div` quality
                       in SV.map (scale qq)
      | otherwise    = SV.map (scale q)
            where q = 200 - quality * 2
                  scale coeff i = fromIntegral . min 255 
                                               . max 1 
                                               $ fromIntegral i * coeff `div` 100
  
  huffmanPackedDecode :: HuffmanPackedTree -> BoolReader s Word8
  huffmanPackedDecode table = getNextBitJpg >>= aux 0
    where aux idx b
              | (v .&. 0x8000) /= 0 = return  0
              | (v .&. 0x4000) /= 0 = return . fromIntegral $ v .&. 0xFF
              | otherwise = getNextBitJpg >>= aux v
            where tableIndex | b = idx + 1
                             | otherwise = idx
                  v = table `SV.unsafeIndex` fromIntegral tableIndex
  
  defaultLumaQuantizationTable :: QuantificationTable
  defaultLumaQuantizationTable = makeMacroBlock
      [16, 11, 10, 16,  24,  40,  51,  61
      ,12, 12, 14, 19,  26,  58,  60,  55
      ,14, 13, 16, 24,  40,  57,  69,  56
      ,14, 17, 22, 29,  51,  87,  80,  62
      ,18, 22, 37, 56,  68, 109, 103,  77
      ,24, 35, 55, 64,  81, 104, 113,  92
      ,49, 64, 78, 87, 103, 121, 120, 101
      ,72, 92, 95, 98, 112, 100, 103,  99
      ]
  
  defaultChromaQuantizationTable :: QuantificationTable
  defaultChromaQuantizationTable = makeMacroBlock
      [17, 18, 24, 47, 99, 99, 99, 99
      ,18, 21, 26, 66, 99, 99, 99, 99
      ,24, 26, 56, 99, 99, 99, 99, 99
      ,47, 66, 99, 99, 99, 99, 99, 99
      ,99, 99, 99, 99, 99, 99, 99, 99
      ,99, 99, 99, 99, 99, 99, 99, 99
      ,99, 99, 99, 99, 99, 99, 99, 99
      ,99, 99, 99, 99, 99, 99, 99, 99
      ]
  
  defaultDcLumaHuffmanTree :: HuffmanTree
  defaultDcLumaHuffmanTree = buildHuffmanTree defaultDcLumaHuffmanTable
  
  -- | From the Table K.3 of ITU-81 (p153)
  defaultDcLumaHuffmanTable :: HuffmanTable
  defaultDcLumaHuffmanTable =
      [ []
      , [0]
      , [1, 2, 3, 4, 5]
      , [6]
      , [7]
      , [8]
      , [9]
      , [10]
      , [11]
      , []
      , []
      , []
      , []
      , []
      , []
      , []
      ]
  
  defaultDcChromaHuffmanTree :: HuffmanTree
  defaultDcChromaHuffmanTree = buildHuffmanTree defaultDcChromaHuffmanTable
  
  -- | From the Table K.4 of ITU-81 (p153)
  defaultDcChromaHuffmanTable :: HuffmanTable
  defaultDcChromaHuffmanTable = 
      [ []
      , [0, 1, 2]
      , [3]
      , [4]
      , [5]
      , [6]
      , [7]
      , [8]
      , [9]
      , [10]
      , [11]
      , []
      , []
      , []
      , []
      , []
      ]
  
  defaultAcLumaHuffmanTree :: HuffmanTree
  defaultAcLumaHuffmanTree = buildHuffmanTree defaultAcLumaHuffmanTable
  
  -- | From the Table K.5 of ITU-81 (p154)
  defaultAcLumaHuffmanTable :: HuffmanTable
  defaultAcLumaHuffmanTable =
      [ []
      , [0x01, 0x02]
      , [0x03]
      , [0x00, 0x04, 0x11]
      , [0x05, 0x12, 0x21]
      , [0x31, 0x41]
      , [0x06, 0x13, 0x51, 0x61]
      , [0x07, 0x22, 0x71]
      , [0x14, 0x32, 0x81, 0x91, 0xA1]
      , [0x08, 0x23, 0x42, 0xB1, 0xC1]
      , [0x15, 0x52, 0xD1, 0xF0]
      , [0x24, 0x33, 0x62, 0x72]
      , []
      , []
      , [0x82]
      , [0x09, 0x0A, 0x16, 0x17, 0x18, 0x19, 0x1A, 0x25, 0x26, 0x27, 0x28, 0x29, 0x2A, 0x34, 0x35
        ,0x36, 0x37, 0x38, 0x39, 0x3A, 0x43, 0x44, 0x45, 0x46, 0x47, 0x48, 0x49, 0x4A, 0x53, 0x54
        ,0x55, 0x56, 0x57, 0x58, 0x59, 0x5A, 0x63, 0x64, 0x65, 0x66, 0x67, 0x68, 0x69, 0x6A, 0x73
        ,0x74, 0x75, 0x76, 0x77, 0x78, 0x79, 0x7A, 0x83, 0x84, 0x85, 0x86, 0x87, 0x88, 0x89, 0x8A
        ,0x92, 0x93, 0x94, 0x95, 0x96, 0x97, 0x98, 0x99, 0x9A, 0xA2, 0xA3, 0xA4, 0xA5, 0xA6, 0xA7
        ,0xA8, 0xA9, 0xAA, 0xB2, 0xB3, 0xB4, 0xB5, 0xB6, 0xB7, 0xB8, 0xB9, 0xBA, 0xC2, 0xC3, 0xC4
        ,0xC5, 0xC6, 0xC7, 0xC8, 0xC9, 0xCA, 0xD2, 0xD3, 0xD4, 0xD5, 0xD6, 0xD7, 0xD8, 0xD9, 0xDA
        ,0xE1, 0xE2, 0xE3, 0xE4, 0xE5, 0xE6, 0xE7, 0xE8, 0xE9, 0xEA, 0xF1, 0xF2, 0xF3, 0xF4, 0xF5
        ,0xF6, 0xF7, 0xF8, 0xF9, 0xFA]
      ]
  
  type HuffmanTable = [[Word8]]
  
  defaultAcChromaHuffmanTree :: HuffmanTree
  defaultAcChromaHuffmanTree = buildHuffmanTree defaultAcChromaHuffmanTable 
  
  defaultAcChromaHuffmanTable :: HuffmanTable
  defaultAcChromaHuffmanTable = 
      [ []
      , [0x00, 0x01]
      , [0x02]
      , [0x03, 0x11]
      , [0x04, 0x05, 0x21, 0x31]
      , [0x06, 0x12, 0x41, 0x51]
      , [0x07, 0x61, 0x71]
      , [0x13, 0x22, 0x32, 0x81]
      , [0x08, 0x14, 0x42, 0x91, 0xA1, 0xB1, 0xC1]
      , [0x09, 0x23, 0x33, 0x52, 0xF0]
      , [0x15, 0x62, 0x72, 0xD1]
      , [0x0A, 0x16, 0x24, 0x34]
      , []
      , [0xE1]
      , [0x25, 0xF1]
      , [ 0x17, 0x18, 0x19, 0x1A, 0x26, 0x27, 0x28, 0x29, 0x2A, 0x35
        , 0x36, 0x37, 0x38, 0x39, 0x3A, 0x43, 0x44, 0x45, 0x46, 0x47
        , 0x48, 0x49, 0x4A, 0x53, 0x54, 0x55, 0x56, 0x57, 0x58, 0x59
        , 0x5A, 0x63, 0x64, 0x65, 0x66, 0x67, 0x68, 0x69, 0x6A, 0x73
        , 0x74, 0x75, 0x76, 0x77, 0x78, 0x79, 0x7A, 0x82, 0x83, 0x84
        , 0x85, 0x86, 0x87, 0x88, 0x89, 0x8A, 0x92, 0x93, 0x94, 0x95
        , 0x96, 0x97, 0x98, 0x99, 0x9A, 0xA2, 0xA3, 0xA4, 0xA5, 0xA6
        , 0xA7, 0xA8, 0xA9, 0xAA, 0xB2, 0xB3, 0xB4, 0xB5, 0xB6, 0xB7
        , 0xB8, 0xB9, 0xBA, 0xC2, 0xC3, 0xC4, 0xC5, 0xC6, 0xC7, 0xC8
        , 0xC9, 0xCA, 0xD2, 0xD3, 0xD4, 0xD5, 0xD6, 0xD7, 0xD8, 0xD9
        , 0xDA, 0xE2, 0xE3, 0xE4, 0xE5, 0xE6, 0xE7, 0xE8, 0xE9, 0xEA
        , 0xF2, 0xF3, 0xF4, 0xF5, 0xF6, 0xF7, 0xF8, 0xF9, 0xFA
        ]
      ]
  
  