{-# LANGUAGE PackageImports, PatternGuards, ExplicitForAll  #-} 

-- | Reading and writing arrays as uncompressed 24 or 32 bit Windows BMP files.
module Data.Array.Repa.IO.BMP
        ( readImageFromBMP
        , writeImageToBMP)
where
import Data.Array.Repa                          as R
import Data.Array.Repa.Unsafe                   as R
import Data.Array.Repa.Repr.ForeignPtr          as R
import Data.Array.Repa.Repr.ByteString          as R
import Data.Vector.Unboxed                      as U
import Prelude                                  as P
import Foreign.ForeignPtr
import Foreign.Marshal.Alloc
import Data.ByteString.Unsafe                   as B
import Codec.BMP
import Data.Word

-- NOTE: We set most of these functions as NOINLINE so it's easier to understand
--       what's going on in the core programs. The top-level IO functions are
--       only called a few times each, so it doesn't matter if they're not
--       worker/wrappered etc.
        
-- Read -------------------------------------------------------------------------------------------
-- | Read RGB components from a BMP file.
readImageFromBMP
        :: FilePath
        -> IO (Either Error (Array U DIM2 (Word8, Word8, Word8)))

{-# NOINLINE readImageFromBMP #-}
readImageFromBMP :: FilePath
-> IO (Either Error (Array U (DIM1 :. Int) (Word8, Word8, Word8)))
readImageFromBMP FilePath
filePath
 = do   Either Error BMP
ebmp    <- FilePath -> IO (Either Error BMP)
readBMP FilePath
filePath

        case Either Error BMP
ebmp of
         Left Error
err       -> Either Error (Array U (DIM1 :. Int) (Word8, Word8, Word8))
-> IO (Either Error (Array U (DIM1 :. Int) (Word8, Word8, Word8)))
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either Error (Array U (DIM1 :. Int) (Word8, Word8, Word8))
 -> IO (Either Error (Array U (DIM1 :. Int) (Word8, Word8, Word8))))
-> Either Error (Array U (DIM1 :. Int) (Word8, Word8, Word8))
-> IO (Either Error (Array U (DIM1 :. Int) (Word8, Word8, Word8)))
forall a b. (a -> b) -> a -> b
$ Error -> Either Error (Array U (DIM1 :. Int) (Word8, Word8, Word8))
forall a b. a -> Either a b
Left Error
err
         Right BMP
bmp      
          -> do Array U (DIM1 :. Int) (Word8, Word8, Word8)
arr     <- BMP -> IO (Array U (DIM1 :. Int) (Word8, Word8, Word8))
forall {m :: * -> *}.
Monad m =>
BMP -> m (Array U (DIM1 :. Int) (Word8, Word8, Word8))
readImageFromBMP' BMP
bmp
                Either Error (Array U (DIM1 :. Int) (Word8, Word8, Word8))
-> IO (Either Error (Array U (DIM1 :. Int) (Word8, Word8, Word8)))
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return  (Either Error (Array U (DIM1 :. Int) (Word8, Word8, Word8))
 -> IO (Either Error (Array U (DIM1 :. Int) (Word8, Word8, Word8))))
-> Either Error (Array U (DIM1 :. Int) (Word8, Word8, Word8))
-> IO (Either Error (Array U (DIM1 :. Int) (Word8, Word8, Word8)))
forall a b. (a -> b) -> a -> b
$ Array U (DIM1 :. Int) (Word8, Word8, Word8)
-> Either Error (Array U (DIM1 :. Int) (Word8, Word8, Word8))
forall a b. b -> Either a b
Right Array U (DIM1 :. Int) (Word8, Word8, Word8)
arr

readImageFromBMP' :: BMP -> m (Array U (DIM1 :. Int) (Word8, Word8, Word8))
readImageFromBMP' BMP
bmp
 = do   let (Int
width, Int
height) = BMP -> (Int, Int)
bmpDimensions BMP
bmp

        let arr :: Array B (DIM1 :. Int) Word8
arr         = (DIM1 :. Int) -> ByteString -> Array B (DIM1 :. Int) Word8
forall sh. sh -> ByteString -> Array B sh Word8
R.fromByteString (Z
Z Z -> Int -> DIM1
forall tail head. tail -> head -> tail :. head
:. Int
height DIM1 -> Int -> DIM1 :. Int
forall tail head. tail -> head -> tail :. head
:. Int
width Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
4)
                        (ByteString -> Array B (DIM1 :. Int) Word8)
-> ByteString -> Array B (DIM1 :. Int) Word8
forall a b. (a -> b) -> a -> b
$ BMP -> ByteString
unpackBMPToRGBA32 BMP
bmp

        let shapeFn :: p -> DIM1 :. Int
shapeFn p
_   = Z
Z Z -> Int -> DIM1
forall tail head. tail -> head -> tail :. head
:. Int
height DIM1 -> Int -> DIM1 :. Int
forall tail head. tail -> head -> tail :. head
:. Int
width

        -- Read out the components into their own arrays, 
        -- skipping the alpha channel.
        Array U (DIM1 :. Int) Word8
vecRed         <- Array D (DIM1 :. Int) Word8 -> m (Array U (DIM1 :. Int) Word8)
forall r1 sh e r2 (m :: * -> *).
(Load r1 sh e, Target r2 e, Source r2 e, Monad m) =>
Array r1 sh e -> m (Array r2 sh e)
computeP 
                        (Array D (DIM1 :. Int) Word8 -> m (Array U (DIM1 :. Int) Word8))
-> Array D (DIM1 :. Int) Word8 -> m (Array U (DIM1 :. Int) Word8)
forall a b. (a -> b) -> a -> b
$ Array B (DIM1 :. Int) Word8
-> ((DIM1 :. Int) -> DIM1 :. Int)
-> (((DIM1 :. Int) -> Word8) -> (DIM1 :. Int) -> Word8)
-> Array D (DIM1 :. Int) Word8
forall r sh sh' a b.
(Source r a, Shape sh) =>
Array r sh a
-> (sh -> sh') -> ((sh -> a) -> sh' -> b) -> Array D sh' b
unsafeTraverse Array B (DIM1 :. Int) Word8
arr (DIM1 :. Int) -> DIM1 :. Int
forall {p}. p -> DIM1 :. Int
shapeFn
                               (\(DIM1 :. Int) -> Word8
get (DIM1
sh :. Int
x) -> (DIM1 :. Int) -> Word8
get (DIM1
sh DIM1 -> Int -> DIM1 :. Int
forall tail head. tail -> head -> tail :. head
:. (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
4)))

        Array U (DIM1 :. Int) Word8
vecGreen       <- Array D (DIM1 :. Int) Word8 -> m (Array U (DIM1 :. Int) Word8)
forall r1 sh e r2 (m :: * -> *).
(Load r1 sh e, Target r2 e, Source r2 e, Monad m) =>
Array r1 sh e -> m (Array r2 sh e)
computeP 
                        (Array D (DIM1 :. Int) Word8 -> m (Array U (DIM1 :. Int) Word8))
-> Array D (DIM1 :. Int) Word8 -> m (Array U (DIM1 :. Int) Word8)
forall a b. (a -> b) -> a -> b
$ Array B (DIM1 :. Int) Word8
-> ((DIM1 :. Int) -> DIM1 :. Int)
-> (((DIM1 :. Int) -> Word8) -> (DIM1 :. Int) -> Word8)
-> Array D (DIM1 :. Int) Word8
forall r sh sh' a b.
(Source r a, Shape sh) =>
Array r sh a
-> (sh -> sh') -> ((sh -> a) -> sh' -> b) -> Array D sh' b
unsafeTraverse Array B (DIM1 :. Int) Word8
arr (DIM1 :. Int) -> DIM1 :. Int
forall {p}. p -> DIM1 :. Int
shapeFn
                                (\(DIM1 :. Int) -> Word8
get (DIM1
sh :. Int
x) -> (DIM1 :. Int) -> Word8
get (DIM1
sh DIM1 -> Int -> DIM1 :. Int
forall tail head. tail -> head -> tail :. head
:. (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
4 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)))

        Array U (DIM1 :. Int) Word8
vecBlue        <- Array D (DIM1 :. Int) Word8 -> m (Array U (DIM1 :. Int) Word8)
forall r1 sh e r2 (m :: * -> *).
(Load r1 sh e, Target r2 e, Source r2 e, Monad m) =>
Array r1 sh e -> m (Array r2 sh e)
computeP
                        (Array D (DIM1 :. Int) Word8 -> m (Array U (DIM1 :. Int) Word8))
-> Array D (DIM1 :. Int) Word8 -> m (Array U (DIM1 :. Int) Word8)
forall a b. (a -> b) -> a -> b
$ Array B (DIM1 :. Int) Word8
-> ((DIM1 :. Int) -> DIM1 :. Int)
-> (((DIM1 :. Int) -> Word8) -> (DIM1 :. Int) -> Word8)
-> Array D (DIM1 :. Int) Word8
forall r sh sh' a b.
(Source r a, Shape sh) =>
Array r sh a
-> (sh -> sh') -> ((sh -> a) -> sh' -> b) -> Array D sh' b
unsafeTraverse Array B (DIM1 :. Int) Word8
arr (DIM1 :. Int) -> DIM1 :. Int
forall {p}. p -> DIM1 :. Int
shapeFn
                                (\(DIM1 :. Int) -> Word8
get (DIM1
sh :. Int
x) -> (DIM1 :. Int) -> Word8
get (DIM1
sh DIM1 -> Int -> DIM1 :. Int
forall tail head. tail -> head -> tail :. head
:. (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
4 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2)))

        -- O(1). zip the components together
        let vecRGB :: Vector (Word8, Word8, Word8)
vecRGB     = Vector Word8
-> Vector Word8 -> Vector Word8 -> Vector (Word8, Word8, Word8)
forall a b c.
(Unbox a, Unbox b, Unbox c) =>
Vector a -> Vector b -> Vector c -> Vector (a, b, c)
U.zip3 (Array U (DIM1 :. Int) Word8 -> Vector Word8
forall sh e. Array U sh e -> Vector e
toUnboxed Array U (DIM1 :. Int) Word8
vecRed)
                                (Array U (DIM1 :. Int) Word8 -> Vector Word8
forall sh e. Array U sh e -> Vector e
toUnboxed Array U (DIM1 :. Int) Word8
vecGreen)
                                (Array U (DIM1 :. Int) Word8 -> Vector Word8
forall sh e. Array U sh e -> Vector e
toUnboxed Array U (DIM1 :. Int) Word8
vecBlue)
        
        Array U (DIM1 :. Int) (Word8, Word8, Word8)
-> m (Array U (DIM1 :. Int) (Word8, Word8, Word8))
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Array U (DIM1 :. Int) (Word8, Word8, Word8)
 -> m (Array U (DIM1 :. Int) (Word8, Word8, Word8)))
-> Array U (DIM1 :. Int) (Word8, Word8, Word8)
-> m (Array U (DIM1 :. Int) (Word8, Word8, Word8))
forall a b. (a -> b) -> a -> b
$ (DIM1 :. Int)
-> Vector (Word8, Word8, Word8)
-> Array U (DIM1 :. Int) (Word8, Word8, Word8)
forall sh e. sh -> Vector e -> Array U sh e
fromUnboxed (Z
Z Z -> Int -> DIM1
forall tail head. tail -> head -> tail :. head
:. Int
height DIM1 -> Int -> DIM1 :. Int
forall tail head. tail -> head -> tail :. head
:. Int
width) Vector (Word8, Word8, Word8)
vecRGB



-- | Write RGB components to a BMP file.
writeImageToBMP
        :: FilePath
        -> Array U DIM2 (Word8, Word8, Word8)
        -> IO ()

{-# NOINLINE writeImageToBMP #-}
writeImageToBMP :: FilePath -> Array U (DIM1 :. Int) (Word8, Word8, Word8) -> IO ()
writeImageToBMP FilePath
fileName Array U (DIM1 :. Int) (Word8, Word8, Word8)
arrRGB
 = do   let sh :: DIM1 :. Int
sh@(Z
Z :. Int
height :. Int
width)
                        = Array U (DIM1 :. Int) (Word8, Word8, Word8) -> DIM1 :. Int
forall sh. Shape sh => Array U sh (Word8, Word8, Word8) -> sh
forall r e sh. (Source r e, Shape sh) => Array r sh e -> sh
extent Array U (DIM1 :. Int) (Word8, Word8, Word8)
arrRGB

        -- O(1). unzip the components
        let (Vector Word8
vecRed, Vector Word8
vecGreen, Vector Word8
vecBlue)
                = Vector (Word8, Word8, Word8)
-> (Vector Word8, Vector Word8, Vector Word8)
forall a b c.
(Unbox a, Unbox b, Unbox c) =>
Vector (a, b, c) -> (Vector a, Vector b, Vector c)
U.unzip3 (Vector (Word8, Word8, Word8)
 -> (Vector Word8, Vector Word8, Vector Word8))
-> Vector (Word8, Word8, Word8)
-> (Vector Word8, Vector Word8, Vector Word8)
forall a b. (a -> b) -> a -> b
$ Array U (DIM1 :. Int) (Word8, Word8, Word8)
-> Vector (Word8, Word8, Word8)
forall sh e. Array U sh e -> Vector e
toUnboxed Array U (DIM1 :. Int) (Word8, Word8, Word8)
arrRGB


        -- Create a bytestring with all the data
        Ptr Word8
ptr     <- Int -> IO (Ptr Word8)
forall a. Int -> IO (Ptr a)
mallocBytes (Int
height Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
width Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
4)
        ForeignPtr Word8
fptr    <- FinalizerPtr Word8 -> Ptr Word8 -> IO (ForeignPtr Word8)
forall a. FinalizerPtr a -> Ptr a -> IO (ForeignPtr a)
newForeignPtr FinalizerPtr Word8
forall a. FinalizerPtr a
finalizerFree Ptr Word8
ptr

        ForeignPtr Word8 -> Array D (DIM1 :. Int) Word8 -> IO ()
forall r1 sh e.
(Load r1 sh e, Storable e) =>
ForeignPtr e -> Array r1 sh e -> IO ()
computeIntoP ForeignPtr Word8
fptr 
         (Array D (DIM1 :. Int) Word8 -> IO ())
-> Array D (DIM1 :. Int) Word8 -> IO ()
forall a b. (a -> b) -> a -> b
$ Array U (DIM1 :. Int) Word8
-> Array U (DIM1 :. Int) Word8
-> Array U (DIM1 :. Int) Word8
-> Array D (DIM1 :. Int) Word8
-> Array D (DIM1 :. Int) Word8
forall sh r1 a r2 r3 r4.
(Shape sh, Source r1 a, Source r2 a, Source r3 a, Source r4 a) =>
Array r1 (sh :. Int) a
-> Array r2 (sh :. Int) a
-> Array r3 (sh :. Int) a
-> Array r4 (sh :. Int) a
-> Array D (sh :. Int) a
interleave4 
                ((DIM1 :. Int) -> Vector Word8 -> Array U (DIM1 :. Int) Word8
forall sh e. sh -> Vector e -> Array U sh e
fromUnboxed DIM1 :. Int
sh Vector Word8
vecRed)
                ((DIM1 :. Int) -> Vector Word8 -> Array U (DIM1 :. Int) Word8
forall sh e. sh -> Vector e -> Array U sh e
fromUnboxed DIM1 :. Int
sh Vector Word8
vecGreen)
                ((DIM1 :. Int) -> Vector Word8 -> Array U (DIM1 :. Int) Word8
forall sh e. sh -> Vector e -> Array U sh e
fromUnboxed DIM1 :. Int
sh Vector Word8
vecBlue)
                ((DIM1 :. Int)
-> ((DIM1 :. Int) -> Word8) -> Array D (DIM1 :. Int) Word8
forall sh a. sh -> (sh -> a) -> Array D sh a
fromFunction DIM1 :. Int
sh (\DIM1 :. Int
_ -> Word8
255))

        -- Pack the data into a BMP file and write it out.
        ForeignPtr Word8 -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
fptr
         ((Ptr Word8 -> IO ()) -> IO ()) -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
ptr' -> do
                ByteString
bs      <- Ptr Word8 -> Int -> IO () -> IO ByteString
unsafePackCStringFinalizer Ptr Word8
ptr' (Int
width Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
height Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
4) (() -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ())
                let bmp :: BMP
bmp = Int -> Int -> ByteString -> BMP
packRGBA32ToBMP Int
width Int
height ByteString
bs
                FilePath -> BMP -> IO ()
writeBMP FilePath
fileName BMP
bmp


{-
-- Normalise --------------------------------------------------------------------------------------
-- | Normalise a matrix to to [0 .. 1], discarding negative values.
--      If the maximum value is 0 then return the array unchanged.
normalisePositive01
        :: (Shape sh, Fractional a, Ord a)
        => Array sh a
        -> Array sh a

{-# INLINE normalisePositive01 #-}
normalisePositive01 arr 
 = let  mx              = foldAll max 0 arr
        elemFn x
         | x >= 0       = x / mx
         | otherwise    = x
   in   mx `seq`
         if mx == 0 
          then arr
          else R.map elemFn arr

-}