{- | This module provides nessesary abstraction over layers of neural networks.

'AbstractLayer' typeclass defines interface of all layers of neural network model.
Its implementation is probably the most low-leveled abstraction of the "Synapse" library.
Notes on how to correctly implement that typeclass are in the docs for it.

'Layer' is the existential datatype that wraps any 'AbstractLayer' instance.
That is the building block of any neural network.
-}

-- 'TypeFamilies' are needed to instantiate 'DType'.

{-# LANGUAGE TypeFamilies #-}

-- 'ExistentialQuantification' is needed to define 'Layer' datatype.

{-# LANGUAGE ExistentialQuantification #-}


module Synapse.NN.Layers.Layer
    ( -- * 'AbstractLayer' typeclass


      AbstractLayer (inputSize, outputSize, nParameters, getParameters, updateParameters, symbolicForward)
    , forward

      -- * 'Layer' existential datatype


    , Layer (Layer)

      -- * 'LayerConfiguration' type alias

    , LayerConfiguration
    ) where


import Synapse.Tensors (DType)
import Synapse.Tensors.Mat (Mat)

import Synapse.Autograd (Symbolic, SymbolIdentifier, Symbol(unSymbol), SymbolMat, constSymbol)


{- | 'AbstractLayer' typeclass defines basic interface of all layers of neural network model.

Every layer should be able to pass 'Synapse.Autograd.SymbolMat' through itself to produce new 'Synapse.Autograd.SymbolMat'
(make prediction based on its parameters) using 'symbolicForward' function,
which allows for gradients to be calculated after predictions, which in turn makes training possible.

'nParameters', 'getParameters' and 'updateParameters' functions allow training of parameters of the layer.
Their implementations should match - that is 'getParameters' function should return list of length 'nParameters'
and 'updateParameters' should expect a list of the same length with the matrices in the same order as were they in 'getParameters'.

"Synapse" manages gradients and parameters for layers with erased type information using prefix system.
'Synapse.Autograd.SymbolIdentifier' is a prefix for name of symbolic parameters that are used in calculation.
Every used parameter should have unique name to be recognised by the autograd -
it must start with given prefix and end with the numerical index of said parameter.
For example 3rd layer with 2 parameters (weights and bias) should
name its weights symbol \"ml3w1\" and name its bias symbol \"ml3w2\" (\"ml3w\" prefix will be supplied).

Important: this typeclass correct implementation is very important (as it is the \'heart\' of "Synapse" library)
for work of the neural network and training, read the docs thoroughly to ensure that all the invariants are met.
-}
class AbstractLayer l where
    -- | Returns the size of the input that is supported for 'forward' and 'symbolicForward' functions. 'Nothing' means size independence (activation functions are the example).

    inputSize :: l a -> Maybe Int
    -- | Returns the size of the output that is supported for 'forward' and 'symbolicForward' functions. 'Nothing' means size independence (activation functions are the example).

    outputSize :: l a -> Maybe Int
    -- | Returns the number of parameters of this layer.

    nParameters :: l a -> Int
    
    -- | Returns a list of all parameters (those must be of the exact same order as they are named by the layer (check 'symbolicForward' docs)).

    getParameters :: SymbolIdentifier -> l a -> [SymbolMat a]
    -- | Updates parameters based on supplied list (length of that list, the order and the form of parameters is EXACTLY the same as those from 'getParameters')

    updateParameters :: l a -> [Mat a] -> l a

    {- | Passes symbolic matrix through to produce new symbolic matrix, while retaining gradients graph.
    Second matrix is a result of application of regularizer on a layer.
    -}
    symbolicForward :: (Symbolic a, Floating a, Ord a) => SymbolIdentifier -> SymbolMat a -> l a -> (SymbolMat a, SymbolMat a)

-- | Passes matrix through to produce new matrix.

forward :: (AbstractLayer l, Symbolic a, Floating a, Ord a) => Mat a -> l a -> Mat a
forward :: forall (l :: * -> *) a.
(AbstractLayer l, Symbolic a, Floating a, Ord a) =>
Mat a -> l a -> Mat a
forward Mat a
input = Symbol (Mat a) -> Mat a
forall a. Symbol a -> a
unSymbol (Symbol (Mat a) -> Mat a)
-> (l a -> Symbol (Mat a)) -> l a -> Mat a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Symbol (Mat a), Symbol (Mat a)) -> Symbol (Mat a)
forall a b. (a, b) -> a
fst ((Symbol (Mat a), Symbol (Mat a)) -> Symbol (Mat a))
-> (l a -> (Symbol (Mat a), Symbol (Mat a)))
-> l a
-> Symbol (Mat a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SymbolIdentifier
-> Symbol (Mat a) -> l a -> (Symbol (Mat a), Symbol (Mat a))
forall a.
(Symbolic a, Floating a, Ord a) =>
SymbolIdentifier
-> SymbolMat a -> l a -> (SymbolMat a, SymbolMat a)
forall (l :: * -> *) a.
(AbstractLayer l, Symbolic a, Floating a, Ord a) =>
SymbolIdentifier
-> SymbolMat a -> l a -> (SymbolMat a, SymbolMat a)
symbolicForward SymbolIdentifier
forall a. Monoid a => a
mempty (Mat a -> Symbol (Mat a)
forall a. a -> Symbol a
constSymbol Mat a
input)


-- | 'Layer' existential datatype wraps anything that implements 'AbstractLayer'.

data Layer a = forall l. (AbstractLayer l) => Layer (l a)

type instance DType (Layer a) = a

instance AbstractLayer Layer where
    inputSize :: forall a. Layer a -> Maybe Int
inputSize (Layer l a
l) = l a -> Maybe Int
forall a. l a -> Maybe Int
forall (l :: * -> *) a. AbstractLayer l => l a -> Maybe Int
inputSize l a
l
    outputSize :: forall a. Layer a -> Maybe Int
outputSize (Layer l a
l) = l a -> Maybe Int
forall a. l a -> Maybe Int
forall (l :: * -> *) a. AbstractLayer l => l a -> Maybe Int
outputSize l a
l

    nParameters :: forall a. Layer a -> Int
nParameters (Layer l a
l) = l a -> Int
forall a. l a -> Int
forall (l :: * -> *) a. AbstractLayer l => l a -> Int
nParameters l a
l
    getParameters :: forall a. SymbolIdentifier -> Layer a -> [SymbolMat a]
getParameters SymbolIdentifier
prefix (Layer l a
l) = SymbolIdentifier -> l a -> [SymbolMat a]
forall a. SymbolIdentifier -> l a -> [SymbolMat a]
forall (l :: * -> *) a.
AbstractLayer l =>
SymbolIdentifier -> l a -> [SymbolMat a]
getParameters SymbolIdentifier
prefix l a
l
    updateParameters :: forall a. Layer a -> [Mat a] -> Layer a
updateParameters (Layer l a
l) = l a -> Layer a
forall a (l :: * -> *). AbstractLayer l => l a -> Layer a
Layer (l a -> Layer a) -> ([Mat a] -> l a) -> [Mat a] -> Layer a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. l a -> [Mat a] -> l a
forall a. l a -> [Mat a] -> l a
forall (l :: * -> *) a. AbstractLayer l => l a -> [Mat a] -> l a
updateParameters l a
l

    symbolicForward :: forall a.
(Symbolic a, Floating a, Ord a) =>
SymbolIdentifier
-> SymbolMat a -> Layer a -> (SymbolMat a, SymbolMat a)
symbolicForward SymbolIdentifier
prefix SymbolMat a
input (Layer l a
l) = SymbolIdentifier
-> SymbolMat a -> l a -> (SymbolMat a, SymbolMat a)
forall a.
(Symbolic a, Floating a, Ord a) =>
SymbolIdentifier
-> SymbolMat a -> l a -> (SymbolMat a, SymbolMat a)
forall (l :: * -> *) a.
(AbstractLayer l, Symbolic a, Floating a, Ord a) =>
SymbolIdentifier
-> SymbolMat a -> l a -> (SymbolMat a, SymbolMat a)
symbolicForward SymbolIdentifier
prefix SymbolMat a
input l a
l


-- | 'LayerConfiguration' type alias represents functions that are able to build layers.

type LayerConfiguration l
    =  Int  -- ^ Output size of previous layer.

    -> l    -- ^ Resulting layer.