{-# LANGUAGE BangPatterns, FlexibleContexts #-}

module Gamgine.Image.PNG.Internal.Filters (defilter_scanlines_arr) where

import Data.Array.Storable
import Data.Array.IO

import Data.Word
import Data.Maybe

import qualified Data.ByteString.Lazy as LB

type Width = Int
type Height = Int

-- | Defilter filtered PNG data (data includes scanlines prepended with filter
--   types).
defilter_scanlines_arr :: (Width,Height) -> Int -> LB.ByteString
                       -> IO (StorableArray (Width,Height) Word8)
defilter_scanlines_arr :: (Int, Int)
-> Int -> ByteString -> IO (StorableArray (Int, Int) Word8)
defilter_scanlines_arr (Int
width,Int
height) Int
bpp ByteString
bs = do
  ([Word8]
slTypes, StorableArray (Int, Int) Word8
imgArr) <- (Int, Int)
-> ByteString -> IO ([Word8], StorableArray (Int, Int) Word8)
imageArray (Int
widthInBytes, Int
height) ByteString
bs
  Int -> Int -> StorableArray (Int, Int) Word8 -> [Word8] -> IO ()
doFilter Int
widthInBytes Int
bpp StorableArray (Int, Int) Word8
imgArr [Word8]
slTypes
  StorableArray (Int, Int) Word8
-> IO (StorableArray (Int, Int) Word8)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return StorableArray (Int, Int) Word8
imgArr
 where
   widthInBytes :: Int
widthInBytes = Int
bppInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
width

-- | create mutable array for defiltering the image and break out scanline types
imageArray :: (Width,Height) -> LB.ByteString
           -> IO ([Word8],StorableArray (Int,Int) Word8)
imageArray :: (Int, Int)
-> ByteString -> IO ([Word8], StorableArray (Int, Int) Word8)
imageArray (Int
width,Int
height) ByteString
bs = do
  StorableArray (Int, Int) Word8
a <- ((Int, Int), (Int, Int))
-> [Word8] -> IO (StorableArray (Int, Int) Word8)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> [e] -> m (a i e)
newListArray ((Int
0,Int
0), (Int
heightInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1, Int
widthInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)) ([Word8] -> IO (StorableArray (Int, Int) Word8))
-> [Word8] -> IO (StorableArray (Int, Int) Word8)
forall a b. (a -> b) -> a -> b
$ ByteString -> [Word8]
LB.unpack ByteString
imageData
  ([Word8], StorableArray (Int, Int) Word8)
-> IO ([Word8], StorableArray (Int, Int) Word8)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ([Word8]
scanlineTypes, StorableArray (Int, Int) Word8
a)
 where
   imageData :: ByteString
imageData = [ByteString] -> ByteString
LB.concat [ByteString]
scanlineData
   ([Word8]
scanlineTypes, [ByteString]
scanlineData) = [(Word8, ByteString)] -> ([Word8], [ByteString])
forall a b. [(a, b)] -> ([a], [b])
unzip [(Word8, ByteString)]
scanlines
   scanlines :: [(Word8, ByteString)]
scanlines = (ByteString -> (Word8, ByteString))
-> [ByteString] -> [(Word8, ByteString)]
forall a b. (a -> b) -> [a] -> [b]
map (Maybe (Word8, ByteString) -> (Word8, ByteString)
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe (Word8, ByteString) -> (Word8, ByteString))
-> (ByteString -> Maybe (Word8, ByteString))
-> ByteString
-> (Word8, ByteString)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Maybe (Word8, ByteString)
LB.uncons) (ByteString -> [ByteString]
chop ByteString
bs)
   chop :: ByteString -> [ByteString]
chop ByteString
b
       | ByteString -> Bool
LB.null ByteString
b   = []
       | Bool
otherwise   = let (ByteString
sl,ByteString
rest) = Int64 -> ByteString -> (ByteString, ByteString)
LB.splitAt Int64
slWidth ByteString
b
                       in ByteString
sl ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
: ByteString -> [ByteString]
chop ByteString
rest
   slWidth :: Int64
slWidth = Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
widthInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)

-- | Perform filtering on the image array (standard adaptive filters supported)
doFilter :: Width -> Int -> StorableArray (Int,Int) Word8 -> [Word8]
         -> IO ()
doFilter :: Int -> Int -> StorableArray (Int, Int) Word8 -> [Word8] -> IO ()
doFilter Int
width Int
bpp StorableArray (Int, Int) Word8
image [Word8]
scanlineTypes = [Word8] -> Int -> IO ()
forall {a}. (Eq a, Num a) => [a] -> Int -> IO ()
doFilter' [Word8]
scanlineTypes Int
0
 where
   doFilter' :: [a] -> Int -> IO ()
doFilter'       [] Int
_ = () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
   doFilter' (a
0:[a]
rest) !Int
y = [a] -> Int -> IO ()
doFilter' [a]
rest (Int
yInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)        -- no filter on this line
   doFilter' (a
1:[a]
rest) !Int
y = Int -> IO ()
sub_filter Int
0 IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> [a] -> Int -> IO ()
doFilter' [a]
rest (Int
yInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
     where sub_filter :: Int -> IO ()
sub_filter !Int
x
               | Int
xInt -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<Int
width = do Word8
subx <- (Int, Int) -> IO Word8
forall {m :: * -> *}.
MArray StorableArray Word8 m =>
(Int, Int) -> m Word8
readByte (Int
y,Int
x)
                              Word8
raw  <- (Int, Int) -> IO Word8
forall {m :: * -> *}.
MArray StorableArray Word8 m =>
(Int, Int) -> m Word8
readByte (Int
y,Int
xInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
bpp)
                              (Int, Int) -> Word8 -> IO ()
writeByte (Int
y,Int
x) (Word8
subxWord8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
+Word8
raw)
                              Int -> IO ()
sub_filter (Int
xInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
               | Bool
otherwise = () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
   doFilter' (a
2:[a]
rest) !Int
y = Int -> IO ()
up_filter Int
0 IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> [a] -> Int -> IO ()
doFilter' [a]
rest (Int
yInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
     where up_filter :: Int -> IO ()
up_filter !Int
x
               | Int
xInt -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<Int
width = do Word8
upx <- (Int, Int) -> IO Word8
forall {m :: * -> *}.
MArray StorableArray Word8 m =>
(Int, Int) -> m Word8
readByte (Int
y,Int
x)
                              Word8
prior <- (Int, Int) -> IO Word8
forall {m :: * -> *}.
MArray StorableArray Word8 m =>
(Int, Int) -> m Word8
readByte (Int
yInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1,Int
x)
                              (Int, Int) -> Word8 -> IO ()
writeByte (Int
y,Int
x) (Word8
upxWord8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
+Word8
prior)
                              Int -> IO ()
up_filter (Int
xInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
               | Bool
otherwise = () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
   doFilter' (a
3:[a]
rest) !Int
y = Int -> IO ()
avg_filter Int
0 IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> [a] -> Int -> IO ()
doFilter' [a]
rest (Int
yInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
     where avg_filter :: Int -> IO ()
avg_filter !Int
x
               | Int
xInt -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<Int
width = do Word8
avgx <- (Int, Int) -> IO Word8
forall {m :: * -> *}.
MArray StorableArray Word8 m =>
(Int, Int) -> m Word8
readByte (Int
y,Int
x)
                              Word8
raw <- (Int, Int) -> IO Word8
forall {m :: * -> *}.
MArray StorableArray Word8 m =>
(Int, Int) -> m Word8
readByte (Int
y,Int
xInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
bpp)
                              Word8
prior <- (Int, Int) -> IO Word8
forall {m :: * -> *}.
MArray StorableArray Word8 m =>
(Int, Int) -> m Word8
readByte (Int
yInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1,Int
x)
                              let s :: Word16
s = ((Word8 -> Word16
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
raw Word16 -> Word16 -> Word16
forall a. Num a => a -> a -> a
+ Word8 -> Word16
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
prior) Word16 -> Word16 -> Word16
forall a. Integral a => a -> a -> a
`div` (Word16
2::Word16))
                              (Int, Int) -> Word8 -> IO ()
writeByte (Int
y,Int
x) (Word8
avgx Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
+ Word16 -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word16
s)
                              Int -> IO ()
avg_filter (Int
xInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
               | Bool
otherwise = () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
   doFilter' (a
4:[a]
rest) !Int
y = Int -> IO ()
paeth_filter Int
0 IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> [a] -> Int -> IO ()
doFilter' [a]
rest (Int
yInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
     where paeth_filter :: Int -> IO ()
paeth_filter !Int
x
               | Int
xInt -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<Int
width = do Word8
paethx <- (Int, Int) -> IO Word8
forall {m :: * -> *}.
MArray StorableArray Word8 m =>
(Int, Int) -> m Word8
readByte (Int
y,Int
x)
                              Word8
a <- (Int, Int) -> IO Word8
forall {m :: * -> *}.
MArray StorableArray Word8 m =>
(Int, Int) -> m Word8
readByte (Int
y,Int
xInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
bpp)
                              Word8
b <- (Int, Int) -> IO Word8
forall {m :: * -> *}.
MArray StorableArray Word8 m =>
(Int, Int) -> m Word8
readByte (Int
yInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1,Int
x)
                              Word8
c <- (Int, Int) -> IO Word8
forall {m :: * -> *}.
MArray StorableArray Word8 m =>
(Int, Int) -> m Word8
readByte (Int
yInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1, Int
xInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
bpp)
                              (Int, Int) -> Word8 -> IO ()
writeByte (Int
y,Int
x) (Word8
paethx Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
+ Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Int -> Int -> Int
paeth_predictor (Word8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
a) (Word8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
b) (Word8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
c)))
                              Int -> IO ()
paeth_filter (Int
xInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
               | Bool
otherwise = () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
   doFilter' (a
_:[a]
rest) !Int
y = [a] -> Int -> IO ()
doFilter' [a]
rest (Int
yInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) -- unknown filter, leave as is (and mess up image ;)
   {-# INLINE readByte #-}
   readByte :: (Int, Int) -> m Word8
readByte (!Int
y,!Int
x) = if Int
xInt -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<Int
0 then Word8 -> m Word8
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Word8
0 else StorableArray (Int, Int) Word8 -> (Int, Int) -> m Word8
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray StorableArray (Int, Int) Word8
image (Int
y,Int
x)
   {-# INLINE writeByte #-}
   writeByte :: (Int, Int) -> Word8 -> IO ()
writeByte = StorableArray (Int, Int) Word8 -> (Int, Int) -> Word8 -> IO ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray StorableArray (Int, Int) Word8
image

{-# INLINE paeth_predictor #-}
paeth_predictor :: Int -> Int -> Int -> Int
paeth_predictor :: Int -> Int -> Int -> Int
paeth_predictor !Int
a !Int
b !Int
c
    | Int
pa Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
pb Bool -> Bool -> Bool
&& Int
pa Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
pc  = Int
a
    | Int
pb Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
pc              = Int
b
    | Bool
otherwise             = Int
c
 where
   p :: Int
p = Int
a Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
b Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
c
   pa :: Int
pa = Int -> Int
forall a. Num a => a -> a
abs(Int
pInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
a)
   pb :: Int
pb = Int -> Int
forall a. Num a => a -> a
abs(Int
pInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
b)
   pc :: Int
pc = Int -> Int
forall a. Num a => a -> a
abs(Int
pInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
c)