{-# LANGUAGE BangPatterns, FlexibleContexts, MultiParamTypeClasses
           , TypeFamilies #-}

module Vision.Image.Class (
    -- * Classes
      Pixel (..), MaskedImage (..), Image (..), ImageChannel, FromFunction (..)
    , FunctorImage (..)
    -- * Functions
    , (!), (!?), nChannels, pixel
    ) where

import Data.Int
import Data.Vector.Storable (Vector, generate, unfoldr)
import Data.Word
import Foreign.Storable (Storable)
import Prelude hiding (map, read)

import Vision.Primitive (
      Z (..), (:.) (..), Point, Size
    , fromLinearIndex, toLinearIndex, shapeLength
    )

-- Classes ---------------------------------------------------------------------

-- | Determines the number of channels and the type of each pixel of the image
-- and how images are represented.
class Pixel p where
    type PixelChannel p

    -- | Returns the number of channels of the pixel.
    --
    -- Must not consume 'p' (could be 'undefined').
    pixNChannels :: p -> Int

    pixIndex :: p -> Int -> PixelChannel p

instance Pixel Int16 where
    type PixelChannel Int16 = Int16
    pixNChannels :: Int16 -> Int
pixNChannels Int16
_   = Int
1
    pixIndex :: Int16 -> Int -> PixelChannel Int16
pixIndex     Int16
p Int
_ = Int16
p

instance Pixel Int32 where
    type PixelChannel Int32 = Int32
    pixNChannels :: Int32 -> Int
pixNChannels Int32
_   = Int
1
    pixIndex :: Int32 -> Int -> PixelChannel Int32
pixIndex     Int32
p Int
_ = Int32
p

instance Pixel Int where
    type PixelChannel Int = Int
    pixNChannels :: Int -> Int
pixNChannels Int
_   = Int
1
    pixIndex :: Int -> Int -> PixelChannel Int
pixIndex     Int
p Int
_ = Int
p

instance Pixel Word8 where
    type PixelChannel Word8 = Word8
    pixNChannels :: Word8 -> Int
pixNChannels Word8
_   = Int
1
    pixIndex :: Word8 -> Int -> PixelChannel Word8
pixIndex     Word8
p Int
_ = Word8
p

instance Pixel Word16 where
    type PixelChannel Word16 = Word16
    pixNChannels :: Word16 -> Int
pixNChannels Word16
_   = Int
1
    pixIndex :: Word16 -> Int -> PixelChannel Word16
pixIndex     Word16
p Int
_ = Word16
p

instance Pixel Word32 where
    type PixelChannel Word32 = Word32
    pixNChannels :: Word32 -> Int
pixNChannels Word32
_   = Int
1
    pixIndex :: Word32 -> Int -> PixelChannel Word32
pixIndex     Word32
p Int
_ = Word32
p

instance Pixel Word where
    type PixelChannel Word = Word
    pixNChannels :: Word -> Int
pixNChannels Word
_   = Int
1
    pixIndex :: Word -> Int -> PixelChannel Word
pixIndex     Word
p Int
_ = Word
p

instance Pixel Float where
    type PixelChannel Float = Float
    pixNChannels :: Float -> Int
pixNChannels Float
_   = Int
1
    pixIndex :: Float -> Int -> PixelChannel Float
pixIndex     Float
p Int
_ = Float
p

instance Pixel Double where
    type PixelChannel Double = Double
    pixNChannels :: Double -> Int
pixNChannels Double
_   = Int
1
    pixIndex :: Double -> Int -> PixelChannel Double
pixIndex     Double
p Int
_ = Double
p

instance Pixel Bool where
    type PixelChannel Bool = Bool
    pixNChannels :: Bool -> Int
pixNChannels Bool
_   = Int
1
    pixIndex :: Bool -> Int -> PixelChannel Bool
pixIndex     Bool
p Int
_ = Bool
p

-- | Provides an abstraction for images which are not defined for each of their
-- pixels. The interface is similar to 'Image' except that indexing functions
-- don't always return.
--
-- Image origin (@'ix2' 0 0@) is located in the upper left corner.
class Storable (ImagePixel i) => MaskedImage i where
    type ImagePixel i

    shape :: i -> Size

    -- | Returns the pixel\'s value at 'Z :. y, :. x'.
    maskedIndex :: i -> Point -> Maybe (ImagePixel i)
    maskedIndex i
img = (i
img forall i. MaskedImage i => i -> Int -> Maybe (ImagePixel i)
`maskedLinearIndex`) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall sh. Shape sh => sh -> sh -> Int
toLinearIndex (forall i. MaskedImage i => i -> Point
shape i
img)
    {-# INLINE maskedIndex #-}

    -- | Returns the pixel\'s value as if the image was a single dimension
    -- vector (row-major representation).
    maskedLinearIndex :: i -> Int -> Maybe (ImagePixel i)
    maskedLinearIndex i
img = (i
img forall i. MaskedImage i => i -> Point -> Maybe (ImagePixel i)
`maskedIndex`) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall sh. Shape sh => sh -> Int -> sh
fromLinearIndex (forall i. MaskedImage i => i -> Point
shape i
img)
    {-# INLINE maskedLinearIndex #-}

    -- | Returns the non-masked values of the image.
    values :: i -> Vector (ImagePixel i)
    values !i
img =
        forall a b. Storable a => (b -> Maybe (a, b)) -> b -> Vector a
unfoldr Int -> Maybe (ImagePixel i, Int)
step Int
0
      where
        !n :: Int
n = forall sh. Shape sh => sh -> Int
shapeLength (forall i. MaskedImage i => i -> Point
shape i
img)

        step :: Int -> Maybe (ImagePixel i, Int)
step !Int
i | Int
i forall a. Ord a => a -> a -> Bool
>= Int
n                              = forall a. Maybe a
Nothing
                | Just ImagePixel i
p <- i
img forall i. MaskedImage i => i -> Int -> Maybe (ImagePixel i)
`maskedLinearIndex` Int
i = forall a. a -> Maybe a
Just (ImagePixel i
p, Int
i forall a. Num a => a -> a -> a
+ Int
1)
                | Bool
otherwise                           = Int -> Maybe (ImagePixel i, Int)
step (Int
i forall a. Num a => a -> a -> a
+ Int
1)
    {-# INLINE values #-}

    {-# MINIMAL shape, (maskedIndex | maskedLinearIndex) #-}

type ImageChannel i = PixelChannel (ImagePixel i)

-- | Provides an abstraction over the internal representation of an image.
--
-- Image origin (@'ix2' 0 0@) is located in the upper left corner.
class MaskedImage i => Image i where
    -- | Returns the pixel value at 'Z :. y :. x'.
    index :: i -> Point -> ImagePixel i
    index i
img = (i
img forall i. Image i => i -> Int -> ImagePixel i
`linearIndex`) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall sh. Shape sh => sh -> sh -> Int
toLinearIndex (forall i. MaskedImage i => i -> Point
shape i
img)
    {-# INLINE index #-}

    -- | Returns the pixel value as if the image was a single dimension vector
    -- (row-major representation).
    linearIndex :: i -> Int -> ImagePixel i
    linearIndex i
img = (i
img forall i. Image i => i -> Point -> ImagePixel i
`index`) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall sh. Shape sh => sh -> Int -> sh
fromLinearIndex (forall i. MaskedImage i => i -> Point
shape i
img)
    {-# INLINE linearIndex #-}

    -- | Returns every pixel values as if the image was a single dimension
    -- vector (row-major representation).
    vector :: i -> Vector (ImagePixel i)
    vector i
img = forall a. Storable a => Int -> (Int -> a) -> Vector a
generate (forall sh. Shape sh => sh -> Int
shapeLength forall a b. (a -> b) -> a -> b
$ forall i. MaskedImage i => i -> Point
shape i
img) (i
img forall i. Image i => i -> Int -> ImagePixel i
`linearIndex`)
    {-# INLINE vector #-}

    {-# MINIMAL index | linearIndex #-}

-- | Provides ways to construct an image from a function.
class FromFunction i where
    type FromFunctionPixel i

    -- | Generates an image by calling the given function for each pixel of the
    -- constructed image.
    fromFunction :: Size -> (Point -> FromFunctionPixel i) -> i

    -- | Generates an image by calling the last function for each pixel of the
    -- constructed image.
    --
    -- The first function is called for each line, generating a line invariant
    -- value.
    --
    -- This function is faster for some image representations as some recurring
    -- computation can be cached.
    fromFunctionLine :: Size -> (Int -> a)
                     -> (a -> Point -> FromFunctionPixel i) -> i
    fromFunctionLine Point
size Int -> a
line a -> Point -> FromFunctionPixel i
f =
        forall i.
FromFunction i =>
Point -> (Point -> FromFunctionPixel i) -> i
fromFunction Point
size (\pt :: Point
pt@(DIM0
Z :. Int
y :. Int
_) -> a -> Point -> FromFunctionPixel i
f (Int -> a
line Int
y) Point
pt)
    {-# INLINE fromFunctionLine #-}

    -- | Generates an image by calling the last function for each pixel of the
    -- constructed image.
    --
    -- The first function is called for each column, generating a column
    -- invariant value.
    --
    -- This function *can* be faster for some image representations as some
    -- recurring computations can be cached. However, it may requires a vector
    -- allocation for these values. If the column invariant is cheap to
    -- compute, prefer 'fromFunction'.
    fromFunctionCol :: Storable b => Size -> (Int -> b)
                    -> (b -> Point -> FromFunctionPixel i) -> i
    fromFunctionCol Point
size Int -> b
col b -> Point -> FromFunctionPixel i
f =
        forall i.
FromFunction i =>
Point -> (Point -> FromFunctionPixel i) -> i
fromFunction Point
size (\pt :: Point
pt@(DIM0
Z :. Int
_ :. Int
x) -> b -> Point -> FromFunctionPixel i
f (Int -> b
col Int
x) Point
pt)
    {-# INLINE fromFunctionCol #-}

    -- | Generates an image by calling the last function for each pixel of the
    -- constructed image.
    --
    -- The two first functions are called for each line and for each column,
    -- respectively, generating common line and column invariant values.
    --
    -- This function is faster for some image representations as some recurring
    -- computation can be cached. However, it may requires a vector
    -- allocation for column values. If the column invariant is cheap to
    -- compute, prefer 'fromFunctionLine'.
    fromFunctionCached :: Storable b => Size
                       -> (Int -> a)               -- ^ Line function
                       -> (Int -> b)               -- ^ Column function
                       -> (a -> b -> Point
                           -> FromFunctionPixel i) -- ^ Pixel function
                       -> i
    fromFunctionCached Point
size Int -> a
line Int -> b
col a -> b -> Point -> FromFunctionPixel i
f =
        forall i.
FromFunction i =>
Point -> (Point -> FromFunctionPixel i) -> i
fromFunction Point
size (\pt :: Point
pt@(DIM0
Z :. Int
y :. Int
x) -> a -> b -> Point -> FromFunctionPixel i
f (Int -> a
line Int
y) (Int -> b
col Int
x) Point
pt)
    {-# INLINE fromFunctionCached #-}

    {-# MINIMAL fromFunction #-}

-- | Defines a class for images on which a function can be applied. The class is
-- different from 'Functor' as there could be some constraints and
-- transformations the pixel and image types.
class (MaskedImage src, MaskedImage res) => FunctorImage src res where
    map :: (ImagePixel src -> ImagePixel res) -> src -> res

-- Functions -------------------------------------------------------------------

-- | Alias of 'maskedIndex'.
(!?) :: MaskedImage i => i -> Point -> Maybe (ImagePixel i)
!? :: forall i. MaskedImage i => i -> Point -> Maybe (ImagePixel i)
(!?) = forall i. MaskedImage i => i -> Point -> Maybe (ImagePixel i)
maskedIndex
{-# INLINE (!?) #-}

-- | Alias of 'index'.
(!) :: Image i => i -> Point -> ImagePixel i
! :: forall i. Image i => i -> Point -> ImagePixel i
(!) = forall i. Image i => i -> Point -> ImagePixel i
index
{-# INLINE (!) #-}

-- | Returns the number of channels of an image.
nChannels :: (Pixel (ImagePixel i), MaskedImage i) => i -> Int
nChannels :: forall i. (Pixel (ImagePixel i), MaskedImage i) => i -> Int
nChannels i
img = forall p. Pixel p => p -> Int
pixNChannels (forall i. MaskedImage i => i -> ImagePixel i
pixel i
img)
{-# INLINE nChannels #-}

-- | Returns an 'undefined' instance of a pixel of the image. This is sometime
-- useful to satisfy the type checker as in a call to 'pixNChannels' :
--
-- > nChannels img = pixNChannels (pixel img)
pixel :: MaskedImage i => i -> ImagePixel i
pixel :: forall i. MaskedImage i => i -> ImagePixel i
pixel i
_ = forall a. HasCallStack => a
undefined