content: {-# OPTIONS_HADDOCK show-extensions #-}
  
  {-# LANGUAGE TypeOperators #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  
  {-|
  Module      : Numeric.Neural.Convolution
  Description : convolutional layers
  Copyright   : (c) Lars Brünjes, 2016
  License     : MIT
  Maintainer  : brunjlar@gmail.com
  Stability   : experimental
  Portability : portable
  
  This module defines /convolutional/ layers.
  -}
  
  module Numeric.Neural.Convolution
      ( focus
      , focus'
      , cover
      , cover'
      , convolution
      ) where
  
  import Control.Category
  import Data.FixedSize
  import Data.Functor.Compose   (Compose(..))
  import Data.Proxy
  import Data.Utils
  import GHC.TypeLits
  import GHC.TypeLits.Witnesses
  import Numeric.Neural.Layer
  import Numeric.Neural.Model
  import Prelude                hiding (id, (.))
  
  -- | Focuses on a specific region of a volume.
  --
  focus :: (KnownNat m, KnownNat n, KnownNat d, KnownNat m', KnownNat n')
           => a                -- ^ the element to use for out-of-bound indices
           -> Volume m n d a   -- ^ the original 'Volume'
           -> (Int, Int)       -- ^ the upper left corner of the focused region
           -> Volume m' n' d a
  focus x v (i, j) = generate $ \(i', j', k') -> maybe x id $ v !? (i + i', j + j', k')
  
  -- | Specialization of 'focus' to 'Volume's with numeric values,
  --   padding out-of-bound indices with zero.
  --
  -- >>> :set -XDataKinds
  -- >>> let v = generate (\(i, j, k) -> i + j + k) :: Volume 2 2 3 Int
  -- >>> focus' v (0, 1) :: Volume 1 2 3 Int
  -- Volume (Matrix [[[1,2,3],[0,0,0]]])
  --
  focus' :: (KnownNat m, KnownNat n, KnownNat d, KnownNat m', KnownNat n', Num a)
            => Volume m n d a   -- ^ the original 'Volume'
            -> (Int, Int)       -- ^ the upper left corner of the focused region
            -> Volume m' n' d a
  focus' = focus 0
  
  -- | Covers a 'Volume' with (smaller) regions.
  --
  cover :: forall m n d a s m' n'.
           (KnownNat m, KnownNat n, KnownNat d, KnownNat s, KnownNat m', KnownNat n')
           => a              -- ^ the element to use for out-of-bound indices
           -> Proxy s        -- ^ a proxy to the region size
           -> Int            -- ^ the stride
           -> Volume m n d a -- ^ the original 'Volume'
           -> Matrix m' n' (Vector (s * s * d) a)
  cover x ps stride v = withNatOp (%*) ps ps $
                        withNatOp (%*) (Proxy :: Proxy (s * s)) (Proxy :: Proxy d) $
                        generate $ toVector . volume
  
    where volume :: (Int, Int) -> Volume s s d a
          volume (i, j) = focus x v (stride * i, stride * j)
  
  -- | Specialization of 'cover' to 'Volume's with numeric values.
  --
  -- >>> :set -XDataKinds
  -- >>> let v = generate (\(i, j, _) -> i + j) :: Volume 4 4 1 Int
  -- >>> cover' (Proxy :: Proxy 3) 1 v :: Matrix 2 2 (Vector 9 Int)
  -- Matrix [[[0,1,2,1,2,3,2,3,4],[1,2,3,2,3,4,3,4,5]],[[1,2,3,2,3,4,3,4,5],[2,3,4,3,4,5,4,5,6]]]
  --
  cover' :: forall m n d a s m' n'.
           (KnownNat m, KnownNat n, KnownNat d, KnownNat s, KnownNat m', KnownNat n', Num a)
           => Proxy s        -- ^ a proxy to the region size
           -> Int            -- ^ the stride
           -> Volume m n d a -- ^ the original 'Volume'
           -> Matrix m' n' (Vector (s * s * d) a)
  cover' = cover 0
  
  -- | Convolves a 'Layer' over a 'Volume'.
  --
  convolution :: forall s m n d m' n' d'.
                 (KnownNat s, KnownNat m, KnownNat n, KnownNat d, KnownNat m', KnownNat n')
                 => Proxy s              -- ^ a proxy to the region size
                 -> Int                  -- ^ the stride
                 -> Layer (s * s * d) d' -- ^ the layer to convolve
                 -> Component (Volume m n d) (Volume m' n' d')
  convolution ps stride l = cArr (Diff $ toVolume . unCompose) .
                            cConvolve l .
                            cArr (Diff $ Compose . cover' ps stride)
  
    where
  
      unCompose :: Compose f g a -> f (g a)
      unCompose (Compose x) = x
  
      _ = natVal (Proxy :: Proxy d)
  