content: {-# LANGUAGE BangPatterns #-}
  {-# LANGUAGE CPP #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  
  -- |
  -- Module      : Data.Massiv.Array.Stencil.Internal
  -- Copyright   : (c) Alexey Kuleshevich 2018-2022
  -- License     : BSD3
  -- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
  -- Stability   : experimental
  -- Portability : non-portable
  module Data.Massiv.Array.Stencil.Internal (
    Stencil (..),
    dimapStencil,
    lmapStencil,
    rmapStencil,
  ) where
  
  import Control.Applicative
  import Control.DeepSeq
  import Data.Massiv.Core.Common
  
  -- | Stencil is abstract description of how to handle elements in the neighborhood of
  -- every array cell in order to compute a value for the cells in the new array. Use
  -- `Data.Massiv.Array.makeStencil` and `Data.Massiv.Array.makeConvolutionStencil` in order
  -- to create a stencil.
  data Stencil ix e a = Stencil
    { stencilSize :: !(Sz ix)
    , stencilCenter :: !ix
    , stencilFunc :: (ix -> e) -> (ix -> e) -> ix -> a
    }
  
  instance Index ix => NFData (Stencil ix e a) where
    rnf (Stencil sz ix f) = sz `deepseq` ix `deepseq` f `seq` ()
  
  instance Functor (Stencil ix e) where
    fmap = rmapStencil
    {-# INLINE fmap #-}
  
  -- Profunctor
  
  -- | A Profunctor dimap. Same caviat applies as in `lmapStencil`
  --
  -- @since 0.2.3
  dimapStencil :: (c -> d) -> (a -> b) -> Stencil ix d a -> Stencil ix c b
  dimapStencil f g stencil@Stencil{stencilFunc = sf} = stencil{stencilFunc = sf'}
    where
      sf' us s = g . sf (f . us) (f . s)
      {-# INLINE sf' #-}
  {-# INLINE dimapStencil #-}
  
  -- | A contravariant map of a second type parameter. In other words map a function over each element
  -- of the array, that the stencil will be applied to.
  --
  -- __Note__: This map can be very inefficient, since for stencils larger than 1 element in size, the
  -- supllied function will be repeatedly applied to the same element. It is better to simply map that
  -- function over the source array instead.
  --
  -- @since 0.2.3
  lmapStencil :: (c -> d) -> Stencil ix d a -> Stencil ix c a
  lmapStencil f stencil@Stencil{stencilFunc = sf} = stencil{stencilFunc = sf'}
    where
      sf' us s = sf (f . us) (f . s)
      {-# INLINE sf' #-}
  {-# INLINE lmapStencil #-}
  
  -- | A covariant map over the right most type argument. In other words the usual `fmap`
  -- from `Functor`:
  --
  -- > fmap == rmapStencil
  --
  -- @since 0.2.3
  rmapStencil :: (a -> b) -> Stencil ix e a -> Stencil ix e b
  rmapStencil f stencil@Stencil{stencilFunc = sf} = stencil{stencilFunc = sf'}
    where
      sf' us s = f . sf us s
      {-# INLINE sf' #-}
  {-# INLINE rmapStencil #-}
  
  unionStencilCenters :: Index ix => Stencil ix e1 a1 -> Stencil ix e2 a2 -> ix
  unionStencilCenters (Stencil _ sC1 _) (Stencil _ sC2 _) = liftIndex2 max sC1 sC2
  {-# INLINE unionStencilCenters #-}
  
  unionStencilSizes :: Index ix => ix -> Stencil ix e1 a1 -> Stencil ix e2 a2 -> Sz ix
  unionStencilSizes maxCenter (Stencil (SafeSz sSz1) sC1 _) (Stencil (SafeSz sSz2) sC2 _) =
    Sz $ liftIndex2 (+) maxCenter $ liftIndex2 max (liftIndex2 (-) sSz1 sC1) (liftIndex2 (-) sSz2 sC2)
  {-# INLINE unionStencilSizes #-}
  
  -- TODO: Test interchange law (u <*> pure y = pure ($ y) <*> u)
  instance Index ix => Applicative (Stencil ix e) where
    pure a = Stencil oneSz zeroIndex (\_ _ _ -> a)
    {-# INLINE pure #-}
    (<*>) s1@(Stencil _ _ f1) s2@(Stencil _ _ f2) = Stencil newSz maxCenter stF
      where
        stF ug gV !ix = f1 ug gV ix (f2 ug gV ix)
        {-# INLINE stF #-}
        !newSz = unionStencilSizes maxCenter s1 s2
        !maxCenter = unionStencilCenters s1 s2
    {-# INLINE (<*>) #-}
  
  #if MIN_VERSION_base(4,10,0)
    liftA2 f s1@(Stencil _ _ f1) s2@(Stencil _ _ f2) = Stencil newSz maxCenter stF
      where
        stF ug gV !ix = f (f1 ug gV ix) (f2 ug gV ix)
        {-# INLINE stF #-}
        !newSz = unionStencilSizes maxCenter s1 s2
        !maxCenter = unionStencilCenters s1 s2
    {-# INLINE liftA2 #-}
  #endif
  
  instance (Index ix, Num a) => Num (Stencil ix e a) where
    (+) = liftA2 (+)
    {-# INLINE (+) #-}
    (-) = liftA2 (-)
    {-# INLINE (-) #-}
    (*) = liftA2 (*)
    {-# INLINE (*) #-}
    negate = fmap negate
    {-# INLINE negate #-}
    abs = fmap abs
    {-# INLINE abs #-}
    signum = fmap signum
    {-# INLINE signum #-}
    fromInteger = pure . fromInteger
    {-# INLINE fromInteger #-}
  
  instance (Index ix, Fractional a) => Fractional (Stencil ix e a) where
    (/) = liftA2 (/)
    {-# INLINE (/) #-}
    recip = fmap recip
    {-# INLINE recip #-}
    fromRational = pure . fromRational
    {-# INLINE fromRational #-}
  
  instance (Index ix, Floating a) => Floating (Stencil ix e a) where
    pi = pure pi
    {-# INLINE pi #-}
    exp = fmap exp
    {-# INLINE exp #-}
    log = fmap log
    {-# INLINE log #-}
    sqrt = fmap sqrt
    {-# INLINE sqrt #-}
    (**) = liftA2 (**)
    {-# INLINE (**) #-}
    logBase = liftA2 logBase
    {-# INLINE logBase #-}
    sin = fmap sin
    {-# INLINE sin #-}
    cos = fmap cos
    {-# INLINE cos #-}
    tan = fmap tan
    {-# INLINE tan #-}
    asin = fmap asin
    {-# INLINE asin #-}
    acos = fmap acos
    {-# INLINE acos #-}
    atan = fmap atan
    {-# INLINE atan #-}
    sinh = fmap sinh
    {-# INLINE sinh #-}
    cosh = fmap cosh
    {-# INLINE cosh #-}
    tanh = fmap tanh
    {-# INLINE tanh #-}
    asinh = fmap asinh
    {-# INLINE asinh #-}
    acosh = fmap acosh
    {-# INLINE acosh #-}
    atanh = fmap atanh
    {-# INLINE atanh #-}
  