-- | Basic definitions for stencil handling.
module Data.Array.Repa.Stencil.Base
        ( Boundary      (..)
        , Stencil       (..)
        , makeStencil, makeStencil2)
where
import Data.Array.Repa.Index

-- | How to handle the case when the stencil lies partly outside the array.
data Boundary a
        -- | Use a fixed value for border regions.
        = BoundFixed !a

        -- | Treat points outside the array as having a constant value.
        | BoundConst !a

        -- | Clamp points outside to the same value as the edge pixel.
        | BoundClamp
        deriving (Int -> Boundary a -> ShowS
[Boundary a] -> ShowS
Boundary a -> String
(Int -> Boundary a -> ShowS)
-> (Boundary a -> String)
-> ([Boundary a] -> ShowS)
-> Show (Boundary a)
forall a. Show a => Int -> Boundary a -> ShowS
forall a. Show a => [Boundary a] -> ShowS
forall a. Show a => Boundary a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Boundary a] -> ShowS
$cshowList :: forall a. Show a => [Boundary a] -> ShowS
show :: Boundary a -> String
$cshow :: forall a. Show a => Boundary a -> String
showsPrec :: Int -> Boundary a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> Boundary a -> ShowS
Show)


-- | Represents a convolution stencil that we can apply to array.
--   Only statically known stencils are supported right now.
data Stencil sh a

        -- | Static stencils are used when the coefficients are fixed,
        --   and known at compile time.
        = StencilStatic
        { Stencil sh a -> sh
stencilExtent :: !sh
        , Stencil sh a -> a
stencilZero   :: !a
        , Stencil sh a -> sh -> a -> a -> a
stencilAcc    :: !(sh -> a -> a -> a) }


-- | Make a stencil from a function yielding coefficients at each index.
makeStencil
        :: Num a
        => sh                   -- ^ Extent of stencil.
        -> (sh -> Maybe a)      -- ^ Get the coefficient at this index.
        -> Stencil sh a

{-# INLINE makeStencil #-}
makeStencil :: sh -> (sh -> Maybe a) -> Stencil sh a
makeStencil sh
ex sh -> Maybe a
getCoeff
 = sh -> a -> (sh -> a -> a -> a) -> Stencil sh a
forall sh a. sh -> a -> (sh -> a -> a -> a) -> Stencil sh a
StencilStatic sh
ex a
0
 ((sh -> a -> a -> a) -> Stencil sh a)
-> (sh -> a -> a -> a) -> Stencil sh a
forall a b. (a -> b) -> a -> b
$ \sh
ix a
val a
acc
        -> case sh -> Maybe a
getCoeff sh
ix of
                Maybe a
Nothing         -> a
acc
                Just a
coeff      -> a
acc a -> a -> a
forall a. Num a => a -> a -> a
+ a
val a -> a -> a
forall a. Num a => a -> a -> a
* a
coeff


-- | Wrapper for `makeStencil` that requires a DIM2 stencil.
makeStencil2
        :: Num a
        => Int -> Int           -- ^ extent of stencil
        -> (DIM2 -> Maybe a)    -- ^ Get the coefficient at this index.
        -> Stencil DIM2 a

{-# INLINE makeStencil2 #-}
makeStencil2 :: Int -> Int -> (DIM2 -> Maybe a) -> Stencil DIM2 a
makeStencil2 Int
height Int
width DIM2 -> Maybe a
getCoeff
        = DIM2 -> (DIM2 -> Maybe a) -> Stencil DIM2 a
forall a sh. Num a => sh -> (sh -> Maybe a) -> Stencil sh a
makeStencil (Z
Z Z -> Int -> Z :. Int
forall tail head. tail -> head -> tail :. head
:. Int
height (Z :. Int) -> Int -> DIM2
forall tail head. tail -> head -> tail :. head
:. Int
width) DIM2 -> Maybe a
getCoeff