{-# LANGUAGE CPP             #-}
{-# LANGUAGE DeriveGeneric   #-}
{-# LANGUAGE RecordWildCards #-}
module Data.STBImage.Immutable (Image(..), unsafeCastImage, flipImage, loadImageBytes, writeNChannelPNG, writeNChannelBMP, writeNChannelTGA) where

import           Data.Either
import           Data.List
import qualified Data.Vector.Storable         as V
import qualified Data.Vector.Storable.Mutable as MV
import           Foreign
import           Foreign.C.String
import           Foreign.C.Types
import           GHC.Generics

import           Data.STBImage.ColorTypes


-- | 'Image' is the least opinionated reasonable type to represent an image, just a vector of pixel 'Color's (laid out top-to-bottom, left-to-right) and a size.
data Image a = Image { Image a -> Vector a
_pixels :: V.Vector a, Image a -> Int
_width :: Int, Image a -> Int
_height :: Int }
           deriving (Image a -> Image a -> Bool
(Image a -> Image a -> Bool)
-> (Image a -> Image a -> Bool) -> Eq (Image a)
forall a. (Storable a, Eq a) => Image a -> Image a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Image a -> Image a -> Bool
$c/= :: forall a. (Storable a, Eq a) => Image a -> Image a -> Bool
== :: Image a -> Image a -> Bool
$c== :: forall a. (Storable a, Eq a) => Image a -> Image a -> Bool
Eq, (forall x. Image a -> Rep (Image a) x)
-> (forall x. Rep (Image a) x -> Image a) -> Generic (Image a)
forall x. Rep (Image a) x -> Image a
forall x. Image a -> Rep (Image a) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall a x. Rep (Image a) x -> Image a
forall a x. Image a -> Rep (Image a) x
$cto :: forall a x. Rep (Image a) x -> Image a
$cfrom :: forall a x. Image a -> Rep (Image a) x
Generic)

instance Show (Image a) where
    show :: Image a -> String
show (Image _ w :: Int
w h :: Int
h) = "Image (" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
w String -> ShowS
forall a. [a] -> [a] -> [a]
++ "x" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
h String -> ShowS
forall a. [a] -> [a] -> [a]
++ ")"

unsafeCastImage :: (Storable a, Storable b) => Image a -> Image b
unsafeCastImage :: Image a -> Image b
unsafeCastImage img :: Image a
img@Image{ _pixels :: forall a. Image a -> Vector a
_pixels = Vector a
_pixels } = Image a
img { _pixels :: Vector b
_pixels = Vector a -> Vector b
forall a b. (Storable a, Storable b) => Vector a -> Vector b
V.unsafeCast Vector a
_pixels }

--

-- | Utility function to flip images, e.g. for use with OpenGL
flipImage :: (Storable a) => Image a -> Image a
flipImage :: Image a -> Image a
flipImage img :: Image a
img@Image{..} = Image a
img { _pixels :: Vector a
_pixels = [Vector a] -> Vector a
forall a. Storable a => [Vector a] -> Vector a
V.concat ([Vector a] -> Vector a)
-> (Vector a -> [Vector a]) -> Vector a -> Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Vector a] -> [Vector a]
forall a. [a] -> [a]
reverse ([Vector a] -> [Vector a])
-> (Vector a -> [Vector a]) -> Vector a -> [Vector a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector a -> [Vector a]
forall a. Storable a => Vector a -> [Vector a]
toRows (Vector a -> Vector a) -> Vector a -> Vector a
forall a b. (a -> b) -> a -> b
$ Vector a
_pixels }
    where
        toRows :: (Storable a) => V.Vector a -> [V.Vector a]
        toRows :: Vector a -> [Vector a]
toRows = (Vector a -> Maybe (Vector a, Vector a)) -> Vector a -> [Vector a]
forall b a. (b -> Maybe (a, b)) -> b -> [a]
unfoldr (\v :: Vector a
v -> if Vector a -> Bool
forall a. Storable a => Vector a -> Bool
V.null Vector a
v then Maybe (Vector a, Vector a)
forall a. Maybe a
Nothing else (Vector a, Vector a) -> Maybe (Vector a, Vector a)
forall a. a -> Maybe a
Just ((Vector a, Vector a) -> Maybe (Vector a, Vector a))
-> (Vector a, Vector a) -> Maybe (Vector a, Vector a)
forall a b. (a -> b) -> a -> b
$ Int -> Vector a -> (Vector a, Vector a)
forall a. Storable a => Int -> Vector a -> (Vector a, Vector a)
V.splitAt Int
_width Vector a
v)
--

foreign import ccall "stb_image.h stbi_load" stbi_load :: CString -> Ptr CInt -> Ptr CInt -> Ptr CInt -> CInt -> IO (Ptr CUChar)
foreign import ccall "stb_image.h stbi_failure_reason" stbi_failure_reason :: IO (CString)
foreign import ccall "stb_image.h &stbi_image_free" stbi_image_free :: FunPtr (Ptr CUChar -> IO ())

loadImageBytes :: Int -> FilePath -> IO (Either String (Image CUChar))
loadImageBytes :: Int -> String -> IO (Either String (Image CUChar))
loadImageBytes comps :: Int
comps path :: String
path = do
    CString
cPath <- String -> IO CString
newCString String
path
    Ptr CInt
widthPtr <- CInt -> IO (Ptr CInt)
forall a. Storable a => a -> IO (Ptr a)
new 0
    Ptr CInt
heightPtr <- CInt -> IO (Ptr CInt)
forall a. Storable a => a -> IO (Ptr a)
new 0
    Ptr CInt
nComponentsPtr <- CInt -> IO (Ptr CInt)
forall a. Storable a => a -> IO (Ptr a)
new 0

    Ptr CUChar
dataPtr <- CString
-> Ptr CInt -> Ptr CInt -> Ptr CInt -> CInt -> IO (Ptr CUChar)
stbi_load CString
cPath Ptr CInt
widthPtr Ptr CInt
heightPtr Ptr CInt
nComponentsPtr (Int -> CInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
comps)

    if Ptr CUChar
dataPtr Ptr CUChar -> Ptr CUChar -> Bool
forall a. Eq a => a -> a -> Bool
/= Ptr CUChar
forall a. Ptr a
nullPtr
        then do
            ForeignPtr CUChar
dataForeignPtr <- FinalizerPtr CUChar -> Ptr CUChar -> IO (ForeignPtr CUChar)
forall a. FinalizerPtr a -> Ptr a -> IO (ForeignPtr a)
newForeignPtr FinalizerPtr CUChar
stbi_image_free Ptr CUChar
dataPtr

            Int
_width  <- CInt -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (CInt -> Int) -> IO CInt -> IO Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ptr CInt -> IO CInt
forall a. Storable a => Ptr a -> IO a
peek Ptr CInt
widthPtr :: IO Int
            Int
_height <- CInt -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (CInt -> Int) -> IO CInt -> IO Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Ptr CInt -> IO CInt
forall a. Storable a => Ptr a -> IO a
peek Ptr CInt
heightPtr :: IO Int

            let _pixels :: Vector CUChar
_pixels = ForeignPtr CUChar -> Int -> Vector CUChar
forall a. Storable a => ForeignPtr a -> Int -> Vector a
V.unsafeFromForeignPtr0 ForeignPtr CUChar
dataForeignPtr (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
comps)

            CString -> IO ()
forall a. Ptr a -> IO ()
free CString
cPath
            Ptr CInt -> IO ()
forall a. Ptr a -> IO ()
free Ptr CInt
widthPtr
            Ptr CInt -> IO ()
forall a. Ptr a -> IO ()
free Ptr CInt
heightPtr
            Ptr CInt -> IO ()
forall a. Ptr a -> IO ()
free Ptr CInt
nComponentsPtr

            Either String (Image CUChar) -> IO (Either String (Image CUChar))
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String (Image CUChar) -> IO (Either String (Image CUChar)))
-> Either String (Image CUChar)
-> IO (Either String (Image CUChar))
forall a b. (a -> b) -> a -> b
$ Image CUChar -> Either String (Image CUChar)
forall a b. b -> Either a b
Right Image :: forall a. Vector a -> Int -> Int -> Image a
Image{..}
        else do
            String
err <- CString -> IO String
peekCString (CString -> IO String) -> IO CString -> IO String
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< IO CString
stbi_failure_reason
            Either String (Image CUChar) -> IO (Either String (Image CUChar))
forall (m :: * -> *) a. Monad m => a -> m a
return (Either String (Image CUChar) -> IO (Either String (Image CUChar)))
-> Either String (Image CUChar)
-> IO (Either String (Image CUChar))
forall a b. (a -> b) -> a -> b
$ String -> Either String (Image CUChar)
forall a b. a -> Either a b
Left String
err

--

foreign import ccall "stb/stb_image_write.h stbi_write_png" stbi_write_png :: CString -> CInt -> CInt -> CInt -> Ptr CUChar -> CInt -> IO CInt
foreign import ccall "stb/stb_image_write.h stbi_write_bmp" stbi_write_bmp :: CString -> CInt -> CInt -> CInt -> Ptr CUChar         -> IO CInt
foreign import ccall "stb/stb_image_write.h stbi_write_tga" stbi_write_tga :: CString -> CInt -> CInt -> CInt -> Ptr CUChar         -> IO CInt

writeNChannelPNG :: (Storable a) => CInt -> FilePath -> Image a -> IO ()
writeNChannelPNG :: CInt -> String -> Image a -> IO ()
writeNChannelPNG comps :: CInt
comps path :: String
path Image{..} = do
    CString
cPath <- String -> IO CString
newCString String
path

    let w :: CInt
w = Int -> CInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
_width :: CInt
    let h :: CInt
h = Int -> CInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
_height :: CInt

    ForeignPtr a -> (Ptr a -> IO CInt) -> IO CInt
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ((ForeignPtr a, Int) -> ForeignPtr a
forall a b. (a, b) -> a
fst ((ForeignPtr a, Int) -> ForeignPtr a)
-> (ForeignPtr a, Int) -> ForeignPtr a
forall a b. (a -> b) -> a -> b
$ Vector a -> (ForeignPtr a, Int)
forall a. Storable a => Vector a -> (ForeignPtr a, Int)
V.unsafeToForeignPtr0 Vector a
_pixels) (\pixBuf :: Ptr a
pixBuf ->
        CString -> CInt -> CInt -> CInt -> Ptr CUChar -> CInt -> IO CInt
stbi_write_png CString
cPath CInt
w CInt
h CInt
comps (Ptr a -> Ptr CUChar
forall a b. Ptr a -> Ptr b
castPtr Ptr a
pixBuf) (CInt
w CInt -> CInt -> CInt
forall a. Num a => a -> a -> a
* CInt
comps) -- bytes per row
        )

    CString -> IO ()
forall a. Ptr a -> IO ()
free CString
cPath

writeNChannelBMP :: (Storable a) => CInt -> FilePath -> Image a -> IO ()
writeNChannelBMP :: CInt -> String -> Image a -> IO ()
writeNChannelBMP comps :: CInt
comps path :: String
path Image{..} = do
    CString
cPath <- String -> IO CString
newCString String
path

    let w :: CInt
w = Int -> CInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
_width :: CInt
    let h :: CInt
h = Int -> CInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
_height :: CInt

    ForeignPtr a -> (Ptr a -> IO CInt) -> IO CInt
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ((ForeignPtr a, Int) -> ForeignPtr a
forall a b. (a, b) -> a
fst ((ForeignPtr a, Int) -> ForeignPtr a)
-> (ForeignPtr a, Int) -> ForeignPtr a
forall a b. (a -> b) -> a -> b
$ Vector a -> (ForeignPtr a, Int)
forall a. Storable a => Vector a -> (ForeignPtr a, Int)
V.unsafeToForeignPtr0 Vector a
_pixels) ((Ptr a -> IO CInt) -> IO CInt) -> (Ptr a -> IO CInt) -> IO CInt
forall a b. (a -> b) -> a -> b
$ CString -> CInt -> CInt -> CInt -> Ptr CUChar -> IO CInt
stbi_write_bmp CString
cPath CInt
w CInt
h CInt
comps (Ptr CUChar -> IO CInt)
-> (Ptr a -> Ptr CUChar) -> Ptr a -> IO CInt
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ptr a -> Ptr CUChar
forall a b. Ptr a -> Ptr b
castPtr

    CString -> IO ()
forall a. Ptr a -> IO ()
free CString
cPath


writeNChannelTGA :: (Storable a) => CInt -> FilePath -> Image a -> IO ()
writeNChannelTGA :: CInt -> String -> Image a -> IO ()
writeNChannelTGA comps :: CInt
comps path :: String
path Image{..} = do
    CString
cPath <- String -> IO CString
newCString String
path

    let w :: CInt
w = Int -> CInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
_width :: CInt
    let h :: CInt
h = Int -> CInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
_height :: CInt

    ForeignPtr a -> (Ptr a -> IO CInt) -> IO CInt
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ((ForeignPtr a, Int) -> ForeignPtr a
forall a b. (a, b) -> a
fst ((ForeignPtr a, Int) -> ForeignPtr a)
-> (ForeignPtr a, Int) -> ForeignPtr a
forall a b. (a -> b) -> a -> b
$ Vector a -> (ForeignPtr a, Int)
forall a. Storable a => Vector a -> (ForeignPtr a, Int)
V.unsafeToForeignPtr0 Vector a
_pixels) ((Ptr a -> IO CInt) -> IO CInt) -> (Ptr a -> IO CInt) -> IO CInt
forall a b. (a -> b) -> a -> b
$ CString -> CInt -> CInt -> CInt -> Ptr CUChar -> IO CInt
stbi_write_tga CString
cPath CInt
w CInt
h CInt
comps (Ptr CUChar -> IO CInt)
-> (Ptr a -> Ptr CUChar) -> Ptr a -> IO CInt
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ptr a -> Ptr CUChar
forall a b. Ptr a -> Ptr b
castPtr

    CString -> IO ()
forall a. Ptr a -> IO ()
free CString
cPath