content: {-# LANGUAGE CPP #-}
  {-# LANGUAGE TypeFamilies #-}
  {-# LANGUAGE TupleSections   #-}
  -- | Module dedicated of Radiance file decompression (.hdr or .pic) file.
  -- Radiance file format is used for High dynamic range imaging.
  module Codec.Picture.HDR( decodeHDR
                          , decodeHDRWithMetadata
                          , encodeHDR
                          , encodeRawHDR
                          , encodeRLENewStyleHDR
                          , writeHDR
                          , writeRLENewStyleHDR
                          ) where
  
  #if !MIN_VERSION_base(4,8,0)
  import Control.Applicative( pure, (<*>), (<$>) )
  #endif
  
  import Data.Bits( Bits, (.&.), (.|.), unsafeShiftL, unsafeShiftR )
  import Data.Char( ord, chr, isDigit )
  import Data.Word( Word8 )
  
  #if !MIN_VERSION_base(4,11,0)
  import Data.Monoid( (<>) )
  #endif
  
  import Control.Monad( when, foldM, foldM_, forM, forM_, unless )
  import Control.Monad.Trans.Class( lift )
  import qualified Data.ByteString as B
  import qualified Data.ByteString.Lazy as L
  import qualified Data.ByteString.Char8 as BC
  
  import Data.List( partition )
  import Data.Binary( Binary( .. ), encode )
  import Data.Binary.Get( Get, getByteString, getWord8 )
  import Data.Binary.Put( putByteString, putLazyByteString )
  
  import Control.Monad.ST( ST, runST )
  import Foreign.Storable ( Storable )
  import Control.Monad.Primitive ( PrimState, PrimMonad )
  import qualified Data.Vector.Storable as V
  import qualified Data.Vector.Storable.Mutable as M
  
  import Codec.Picture.Metadata( Metadatas
                               , SourceFormat( SourceHDR )
                               , basicMetadata )
  import Codec.Picture.InternalHelper
  import Codec.Picture.Types
  import Codec.Picture.VectorByteConversion
  
  #if MIN_VERSION_transformers(0, 4, 0)
  import Control.Monad.Trans.Except( ExceptT, throwE, runExceptT )
  #else
  -- Transfomers 0.3 compat
  import Control.Monad.Trans.Error( Error, ErrorT, throwError, runErrorT )
  
  type ExceptT = ErrorT
  
  throwE :: (Monad m, Error e) => e -> ErrorT e m a
  throwE = throwError
  
  runExceptT :: ErrorT e m a -> m (Either e a)
  runExceptT = runErrorT
  #endif
  
  {-# INLINE (.<<.) #-}
  (.<<.), (.>>.) :: (Bits a) => a -> Int -> a
  (.<<.) = unsafeShiftL
  (.>>.) = unsafeShiftR
  
  {-# INLINE (.<-.) #-}
  (.<-.) :: (PrimMonad m, Storable a)
         => M.STVector (PrimState m) a -> Int -> a -> m ()
  (.<-.) = M.write 
           {-M.unsafeWrite-}
  
  type HDRReader s a = ExceptT String (ST s) a
  
  data RGBE = RGBE !Word8 !Word8 !Word8 !Word8
  
  instance Binary RGBE where
      put (RGBE r g b e) = put r >> put g >> put b >> put e
      get = RGBE <$> get <*> get <*> get <*> get
  
  checkLineLength :: RGBE -> Int
  checkLineLength (RGBE _ _ a b) =
      (fromIntegral a .<<. 8) .|. fromIntegral b
  
  isNewRunLengthMarker :: RGBE -> Bool
  isNewRunLengthMarker (RGBE 2 2 _ _) = True
  isNewRunLengthMarker _ = False
  
  data RadianceFormat =
        FormatRGBE
      | FormatXYZE
  
  radiance32bitRleRGBEFormat, radiance32bitRleXYZEFromat :: B.ByteString
  radiance32bitRleRGBEFormat = BC.pack "32-bit_rle_rgbe"
  radiance32bitRleXYZEFromat = BC.pack "32-bit_rle_xyze"
  
  instance Binary RadianceFormat where
    put FormatRGBE = putByteString radiance32bitRleRGBEFormat
    put FormatXYZE = putByteString radiance32bitRleXYZEFromat
  
    get = getByteString (B.length radiance32bitRleRGBEFormat) >>= format
      where format sig
              | sig == radiance32bitRleRGBEFormat = pure FormatRGBE
              | sig == radiance32bitRleXYZEFromat = pure FormatXYZE
              | otherwise = fail "Unrecognized Radiance format"
  
  toRGBE :: PixelRGBF -> RGBE
  toRGBE (PixelRGBF r g b)
      | d <= 1e-32 = RGBE 0 0 0 0
      | otherwise = RGBE (fix r) (fix g) (fix b) (fromIntegral $ e + 128)
    where d = maximum [r, g, b]
          e = exponent d
          coeff = significand d *  255.9999 / d
          fix v = truncate $ v * coeff
  
  
  dropUntil :: Word8 -> Get ()
  dropUntil c = getWord8 >>= inner
    where inner val | val == c = pure ()
          inner _ = getWord8 >>= inner
  
  getUntil :: (Word8 -> Bool) -> B.ByteString -> Get B.ByteString
  getUntil f initialAcc = getWord8 >>= inner initialAcc
    where inner acc c | f c = pure acc
          inner acc c = getWord8 >>= inner (B.snoc acc c)
  
  data RadianceHeader = RadianceHeader
    { radianceInfos :: [(B.ByteString, B.ByteString)]
    , radianceFormat :: RadianceFormat
    , radianceHeight :: !Int
    , radianceWidth  :: !Int
    , radianceData   :: L.ByteString
    }
  
  radianceFileSignature :: B.ByteString
  radianceFileSignature = BC.pack "#?RADIANCE\n"
  
  unpackColor :: L.ByteString -> Int -> RGBE
  unpackColor str idx = RGBE (at 0) (at 1) (at 2) (at 3)
    where at n = L.index str . fromIntegral $ idx + n
  
  storeColor :: M.STVector s Word8 -> Int -> RGBE -> ST s ()
  storeColor vec idx (RGBE r g b e) = do
      (vec .<-. (idx + 0)) r
      (vec .<-. (idx + 1)) g
      (vec .<-. (idx + 2)) b
      (vec .<-. (idx + 3)) e
  
  parsePair :: Char -> Get (B.ByteString, B.ByteString)
  parsePair firstChar = do
      let eol c = c == fromIntegral (ord '\n')
      line <- getUntil eol B.empty
      case BC.split '=' line of
        [] -> pure (BC.singleton firstChar, B.empty)
        [val] -> pure (BC.singleton firstChar, val)
        [key, val] -> pure (BC.singleton firstChar <> key, val)
        (key : vals) -> pure (BC.singleton firstChar <> key, B.concat vals)
  
  decodeInfos :: Get [(B.ByteString, B.ByteString)]
  decodeInfos = do
      char <- getChar8
      case char of
        -- comment
        '#' -> dropUntil (fromIntegral $ ord '\n') >> decodeInfos
        -- end of header, no more information
        '\n' -> pure []
        -- Classical parsing
        c -> (:) <$> parsePair c <*> decodeInfos
  
  
  -- | Decode an HDR (radiance) image, the resulting image can be:
  --
  --  * 'ImageRGBF'
  --
  decodeHDR :: B.ByteString -> Either String DynamicImage
  decodeHDR = fmap fst . decodeHDRWithMetadata
  
  -- | Equivalent to decodeHDR but with aditional metadatas.
  decodeHDRWithMetadata :: B.ByteString -> Either String (DynamicImage, Metadatas)
  decodeHDRWithMetadata str = runST $ runExceptT $
    case runGet decodeHeader $ L.fromChunks [str] of
      Left err -> throwE err
      Right rez ->
        let meta = basicMetadata SourceHDR (abs $ radianceWidth rez) (abs $ radianceHeight rez) in
        (, meta) . ImageRGBF <$> (decodeRadiancePicture rez >>= lift . unsafeFreezeImage)
  
  getChar8 :: Get Char
  getChar8 = chr . fromIntegral <$> getWord8
  
  isSign :: Char -> Bool
  isSign c = c == '+' || c == '-'
  
  isAxisLetter :: Char -> Bool
  isAxisLetter c = c == 'X' || c == 'Y'
  
  decodeNum :: Get Int
  decodeNum = do
      sign <- getChar8
      letter <- getChar8
      space <- getChar8
  
      unless (isSign sign && isAxisLetter letter && space == ' ')
             (fail "Invalid radiance size declaration")
  
      let numDec acc c | isDigit c =
              getChar8 >>= numDec (acc * 10 + ord c - ord '0')
          numDec acc _
              | sign == '-' = pure $ negate acc
              | otherwise = pure acc
  
      getChar8 >>= numDec 0
  
  copyPrevColor :: M.STVector s Word8 -> Int -> ST s ()
  copyPrevColor scanLine idx = do
      r <- scanLine `M.unsafeRead` (idx - 4)
      g <- scanLine `M.unsafeRead` (idx - 3)
      b <- scanLine `M.unsafeRead` (idx - 2)
      e <- scanLine `M.unsafeRead` (idx - 1)
  
      (scanLine `M.unsafeWrite` (idx + 0)) r
      (scanLine `M.unsafeWrite` (idx + 1)) g
      (scanLine `M.unsafeWrite` (idx + 2)) b
      (scanLine `M.unsafeWrite` (idx + 3)) e
  
  oldStyleRLE :: L.ByteString -> Int -> M.STVector s Word8
              -> HDRReader s Int
  oldStyleRLE inputData initialIdx scanLine = inner initialIdx 0 0
    where maxOutput = M.length scanLine
          maxInput = fromIntegral $ L.length inputData
  
          inner readIdx writeIdx _
              | readIdx >= maxInput || writeIdx >= maxOutput = pure readIdx
          inner readIdx writeIdx shift = do
            let color@(RGBE r g b e) = unpackColor inputData readIdx
                isRun = r == 1 && g == 1 && b == 1
  
            if not isRun
              then do
                lift $ storeColor scanLine writeIdx color
                inner (readIdx + 4) (writeIdx + 4) 0
           
              else do
                let count = fromIntegral e .<<. shift
                lift $ forM_ [0 .. count] $ \i -> copyPrevColor scanLine (writeIdx + 4 * i)
                inner (readIdx + 4) (writeIdx + 4 * count) (shift + 8)
  
  newStyleRLE :: L.ByteString -> Int -> M.STVector s Word8
              -> HDRReader s Int
  newStyleRLE inputData initialIdx scanline = foldM inner initialIdx [0 .. 3]
    where dataAt idx
              | fromIntegral idx >= maxInput = throwE $ "Read index out of bound (" ++ show idx ++ ")"
              | otherwise = pure $ L.index inputData (fromIntegral idx)
  
          maxOutput = M.length scanline
          maxInput = fromIntegral $ L.length inputData
          stride = 4
  
  
          strideSet count destIndex _ | endIndex > maxOutput + stride =
            throwE $ "Out of bound HDR scanline " ++ show endIndex ++ " (max " ++ show maxOutput ++ ")"
              where endIndex = destIndex + count * stride
          strideSet count destIndex val = aux destIndex count
              where aux i 0 =  pure i
                    aux i c = do
                      lift $ (scanline .<-. i) val
                      aux (i + stride) (c - 1)
  
  
          strideCopy _ count destIndex
              | writeEndBound > maxOutput + stride = throwE "Out of bound HDR scanline"
                  where writeEndBound = destIndex + count * stride
          strideCopy sourceIndex count destIndex = aux sourceIndex destIndex count
            where aux _ j 0 = pure j
                  aux i j c = do
                      val <- dataAt i
                      lift $ (scanline .<-. j) val
                      aux (i + 1) (j + stride) (c - 1)
  
          inner readIdx writeIdx
              | readIdx >= maxInput || writeIdx >= maxOutput = pure readIdx
          inner readIdx writeIdx = do
            code <- dataAt readIdx
            if code > 128
              then do
                let repeatCount = fromIntegral code .&. 0x7F
                newVal <- dataAt $ readIdx + 1
                endIndex <- strideSet repeatCount writeIdx newVal
                inner (readIdx + 2) endIndex 
  
              else do
                let iCode = fromIntegral code
                endIndex <- strideCopy (readIdx + 1) iCode writeIdx
                inner (readIdx + iCode + 1) endIndex
  
  instance Binary RadianceHeader where
      get = decodeHeader
      put hdr = do
          putByteString radianceFileSignature
          putByteString $ BC.pack "FORMAT="
          put $ radianceFormat hdr
          let sizeString =
                BC.pack $ "\n\n-Y " ++ show (radianceHeight hdr)
                          ++ " +X " ++ show (radianceWidth hdr) ++ "\n"
          putByteString sizeString
          putLazyByteString $ radianceData hdr
  
  
  decodeHeader :: Get RadianceHeader
  decodeHeader = do
      sig <- getByteString $ B.length radianceFileSignature
      when (sig /= radianceFileSignature)
           (fail "Invalid radiance file signature")
  
      infos <- decodeInfos
      let formatKey = BC.pack "FORMAT"
      case partition (\(k,_) -> k /= formatKey) infos of
        (_, []) -> fail "No radiance format specified"
        (info, [(_, formatString)]) ->
          case runGet get $ L.fromChunks [formatString] of
            Left err -> fail err
            Right format -> do
                (n1, n2, b) <- (,,) <$> decodeNum
                                    <*> decodeNum
                                    <*> getRemainingBytes
                return . RadianceHeader info format n1 n2 $ L.fromChunks [b]
  
        _ -> fail "Multiple radiance format specified"
  
  toFloat :: RGBE -> PixelRGBF
  toFloat (RGBE r g b e) = PixelRGBF rf gf bf
    where f = encodeFloat 1 $ fromIntegral e - (128 + 8)
          rf = (fromIntegral r + 0.0) * f
          gf = (fromIntegral g + 0.0) * f
          bf = (fromIntegral b + 0.0) * f
  
  encodeScanlineColor :: M.STVector s Word8
                      -> M.STVector s Word8
                      -> Int
                      -> ST s Int
  encodeScanlineColor vec outVec outIdx = do
      val <- vec `M.unsafeRead` 0
      runLength 1 0 val 1 outIdx
    where maxIndex = M.length vec
  
          pushRun len val at = do
              (outVec `M.unsafeWrite` at) $ fromIntegral $ len .|. 0x80
              (outVec `M.unsafeWrite` (at + 1)) val
              return $ at + 2
  
          pushData start len at = do
              (outVec `M.unsafeWrite` at) $ fromIntegral len
              let first = start - len
                  end = start - 1
                  offset = at - first + 1
              forM_ [first .. end] $ \i -> do
                  v <- vec `M.unsafeRead` i
                  (outVec `M.unsafeWrite` (offset + i)) v
  
              return $ at + len + 1
  
          -- End of scanline, empty the thing
          runLength run cpy prev idx at | idx >= maxIndex =
              case (run, cpy) of
                  (0, 0) -> pure at
                  (0, n) -> pushData idx n at
                  (n, 0) -> pushRun n prev at
                  (_, _) -> error "HDR - Run length algorithm is wrong"
  
          -- full runlength, we must write the packet
          runLength r@127   _ prev idx at = do
              val <- vec `M.unsafeRead` idx
              pushRun r prev at >>=
                  runLength 1 0 val (idx + 1)
  
          -- full copy, we must write the packet
          runLength   _ c@127    _ idx at = do
              val <- vec `M.unsafeRead` idx
              pushData idx c at >>=
                  runLength 1 0 val (idx + 1)
  
          runLength n 0 prev idx at = do
              val <- vec `M.unsafeRead` idx
              case val == prev of
                 True -> runLength (n + 1) 0 prev (idx + 1) at
                 False | n < 4 -> runLength 0 (n + 1) val (idx + 1) at
                 False ->
                      pushRun n prev at >>=
                          runLength 1 0 val (idx + 1)
  
          runLength 0 n prev idx at = do
              val <- vec `M.unsafeRead` idx
              if val /= prev
                 then runLength 0 (n + 1) val (idx + 1) at
                 else
                  pushData (idx - 1) (n - 1) at >>=
                      runLength (2 :: Int) 0 val (idx + 1)
  
          runLength _ _ _ _ _ =
              error "HDR RLE inconsistent state"
  
  -- | Write an High dynamic range image into a radiance
  -- image file on disk.
  writeHDR :: FilePath -> Image PixelRGBF -> IO ()
  writeHDR filename img = L.writeFile filename $ encodeHDR img
  
  -- | Write a RLE encoded High dynamic range image into a radiance
  -- image file on disk.
  writeRLENewStyleHDR :: FilePath -> Image PixelRGBF -> IO ()
  writeRLENewStyleHDR filename img =
      L.writeFile filename $ encodeRLENewStyleHDR img
  
  -- | Encode an High dynamic range image into a radiance image
  -- file format.
  -- Alias for encodeRawHDR
  encodeHDR :: Image PixelRGBF -> L.ByteString
  encodeHDR = encodeRawHDR
  
  -- | Encode an High dynamic range image into a radiance image
  -- file format. without compression
  encodeRawHDR :: Image PixelRGBF -> L.ByteString
  encodeRawHDR pic = encode descriptor
    where
      newImage = pixelMap rgbeInRgba pic
      -- we are cheating to death here, the layout we want
      -- correspond to the layout of pixelRGBA8, so we
      -- convert
      rgbeInRgba pixel = PixelRGBA8 r g b e
        where RGBE r g b e = toRGBE pixel
  
      descriptor = RadianceHeader
          { radianceInfos = []
          , radianceFormat = FormatRGBE
          , radianceHeight = imageHeight pic
          , radianceWidth  = imageWidth pic
          , radianceData = L.fromChunks [toByteString $ imageData newImage]
          }
  
  
  -- | Encode an High dynamic range image into a radiance image
  -- file format using a light RLE compression. Some problems
  -- seem to arise with some image viewer.
  encodeRLENewStyleHDR :: Image PixelRGBF -> L.ByteString
  encodeRLENewStyleHDR pic = encode $ runST $ do
      let w = imageWidth pic
          h = imageHeight pic
  
      scanLineR <- M.new w :: ST s (M.STVector s Word8)
      scanLineG <- M.new w
      scanLineB <- M.new w
      scanLineE <- M.new w
  
      encoded <-
          forM [0 .. h - 1] $ \line -> do
              buff <- M.new $ w * 4 + w `div` 127 + 2
              let columner col | col >= w = return ()
                  columner col = do
                        let RGBE r g b e = toRGBE $ pixelAt pic col line
                        (scanLineR `M.unsafeWrite` col) r
                        (scanLineG `M.unsafeWrite` col) g
                        (scanLineB `M.unsafeWrite` col) b
                        (scanLineE `M.unsafeWrite` col) e
  
                        columner (col + 1)
  
              columner 0
  
              (buff `M.unsafeWrite` 0) 2
              (buff `M.unsafeWrite` 1) 2
              (buff `M.unsafeWrite` 2) $ fromIntegral ((w .>>. 8) .&. 0xFF)
              (buff `M.unsafeWrite` 3) $ fromIntegral (w .&. 0xFF)
  
              i1 <- encodeScanlineColor scanLineR buff 4        
              i2 <- encodeScanlineColor scanLineG buff i1
              i3 <- encodeScanlineColor scanLineB buff i2
              endIndex <- encodeScanlineColor scanLineE buff i3
  
              (\v -> blitVector v 0 endIndex) <$> V.unsafeFreeze buff
  
      pure RadianceHeader
          { radianceInfos = []
          , radianceFormat = FormatRGBE
          , radianceHeight = h
          , radianceWidth  = w
          , radianceData = L.fromChunks encoded 
          }
      
  
  decodeRadiancePicture :: RadianceHeader -> HDRReader s (MutableImage s PixelRGBF)
  decodeRadiancePicture hdr = do
      let width = abs $ radianceWidth hdr
          height = abs $ radianceHeight hdr
          packedData = radianceData hdr
  
      scanLine <- lift $ M.new $ width * 4
      resultBuffer <- lift $ M.new $ width * height * 3
  
      let scanLineImage = MutableImage
                        { mutableImageWidth = width
                        , mutableImageHeight = 1
                        , mutableImageData = scanLine
                        }
  
          finalImage = MutableImage
                     { mutableImageWidth = width
                     , mutableImageHeight = height
                     , mutableImageData = resultBuffer
                     }
  
      let scanLineExtractor readIdx line = do
            let color = unpackColor packedData readIdx
                inner | isNewRunLengthMarker color = do
                            let calcSize = checkLineLength color
                            when (calcSize /= width)
                                 (throwE "Invalid sanline size")
                            pure $ \idx -> newStyleRLE packedData (idx + 4)
                      | otherwise = pure $ oldStyleRLE packedData
            f <- inner
            newRead <- f readIdx scanLine
            forM_ [0 .. width - 1] $ \i -> do
                -- mokay, it's a hack, but I don't want to define a
                -- pixel instance of RGBE...
                PixelRGBA8 r g b e <- lift $ readPixel scanLineImage i 0
                lift $ writePixel finalImage i line . toFloat $ RGBE r g b e
  
            return newRead
  
      foldM_ scanLineExtractor 0 [0 .. height - 1]
  
      return finalImage
  
  