{-# 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 )
import Data.Monoid( (<>) )
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 pixel
-- type can be :
--
--  * PixelRGBF
--
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