-- | Multilayer perceptrons and backpropagation.
module Goal.Probability.Graphical.NeuralNetwork where


--- Imports ---


-- Goal --

import Goal.Geometry
import Goal.Probability.ExponentialFamily
import Goal.Probability.Graphical

import qualified Data.Vector.Storable as C


--- Neural Networks ---


-- | A mutlilayer perceptron with three layers.
data NeuralNetwork m n o = NeuralNetwork m n o deriving (Eq, Read, Show)


--- Functions ---

splitNeuralNetwork
    :: (Manifold m, Manifold n, Manifold o)
    => Function Mixture Mixture :#: NeuralNetwork m n o
    -> (Natural :#: m, NaturalFunction :#: Tensor m n, Natural :#: n, NaturalFunction :#: Tensor n o)
-- | Splits the 'NeuralNetwork' into its component affine transformations.
splitNeuralNetwork nnp =
    let (NeuralNetwork m n o) = manifold nnp
        tns1 = Tensor m n
        tns2 = Tensor n o
        css = coordinates nnp
        (mcs,css') = C.splitAt (dimension m) css
        (mtx1cs,css'') = C.splitAt (dimension tns1) css'
        (ncs,mtx2cs) = C.splitAt (dimension n) css''
        mp = fromCoordinates m mcs
        mtx1 = fromCoordinates tns1 mtx1cs
        np = fromCoordinates n ncs
        mtx2 = fromCoordinates tns2 mtx2cs
     in (mp,mtx1,np,mtx2)

joinNeuralNetwork
    :: (Manifold m, Manifold n, Manifold o)
    => Natural :#: m
    -> NaturalFunction :#: Tensor m n
    -> Natural :#: n
    -> NaturalFunction :#: Tensor n o
    -> Function Mixture Mixture :#: NeuralNetwork m n o
-- | Construct a 'NeuralNetwork' from component affine transformations.
joinNeuralNetwork mp mtx1 np mtx2 =
    let (Tensor m n) = manifold mtx1
        (Tensor _ o) = manifold mtx2
     in fromCoordinates (NeuralNetwork m n o) $ coordinates mp C.++ coordinates mtx1 C.++ coordinates np C.++ coordinates mtx2

feedForward
    :: (ExponentialFamily m, ExponentialFamily n, Manifold o)
    => Function Mixture Mixture :#: NeuralNetwork m n o
    -> [Mixture :#: o]
    -> ([Natural :#: n], [Mixture :#: n], [Natural :#: m], [Mixture :#: m])
-- | Feeds an input forward through the network, and returns every step of
-- the computation.
feedForward nnp xps =
    let (mp,mtx1,np,mtx2) = splitNeuralNetwork nnp
        nyps = map (<+> np) $ mtx2 >$> xps
        yps = potentialMapping <$> nyps
        nzps = map (<+> mp) $ mtx1 >$> yps
        zps = potentialMapping <$> nzps
     in (nyps,yps,nzps,zps)

feedBackward
    :: (Legendre Natural m, Legendre Natural n, Riemannian Natural m, Riemannian Natural n, Manifold o)
    => Function Mixture Mixture :#: NeuralNetwork m n o
    -> [Mixture :#: o]
    -> [Natural :#: n]
    -> [Mixture :#: n]
    -> [Natural :#: m]
    -> [Natural :#: m]
    -> Differentials :#: Tangent (Function Mixture Mixture) (NeuralNetwork m n o)
-- | Given the results of a feed forward application, back propagates a
-- given error (last input) through the network.
feedBackward nnp xps nyps yps nzps errs1 =
    let (_,mtx1,_,_) = splitNeuralNetwork nnp
        dmps = zipWith legendreFlat nzps errs1
        dmtx1s = [ dmp >.< yp | (dmp,yp) <- zip dmps yps ]
        errs2 = matrixTranspose mtx1 >$> dmps
        dnps = zipWith legendreFlat nyps errs2
        dmtx2s = [ dnp >.< xp | (dnp,xp) <- zip dnps xps ]
     in fromCoordinates (Tangent nnp) $ coordinates (meanPoint dmps) C.++ coordinates (meanPoint dmtx1s)
            C.++ coordinates (meanPoint dnps) C.++ coordinates (meanPoint dmtx2s)

meanSquaredBackpropagation
    :: (Riemannian Natural m, Riemannian Natural n, ExponentialFamily m, ExponentialFamily n, Manifold o)
    => Function Mixture Mixture :#: NeuralNetwork m n o
    -> [Mixture :#: o]
    -> [Mixture :#: m]
    -> Differentials :#: Tangent (Function Mixture Mixture) (NeuralNetwork m n o)
-- | Backpropagation algorithm with the mean squared error function.
meanSquaredBackpropagation nnp xps tps =
    let (nyps,yps,nzps,zps) = feedForward nnp xps
        errs1 = [ alterChart Natural $ zp <-> tp | (tp,zp) <- zip tps zps ]
     in feedBackward nnp xps nyps yps nzps errs1


--- Instances ---


instance (Manifold m, Manifold n, Manifold o) => Manifold (NeuralNetwork m n o) where
    dimension (NeuralNetwork m n o) =  dimension m + dimension m * dimension n + dimension n + dimension n * dimension o

instance (ExponentialFamily m, ExponentialFamily n, Manifold o) => Map (NeuralNetwork m n o) where
    type Domain (NeuralNetwork m n o) = o
    domain (NeuralNetwork _ _ o) = o
    type Codomain (NeuralNetwork m n o) = m
    codomain (NeuralNetwork m _ _) = m

instance (ExponentialFamily m, ExponentialFamily n, Manifold o) => Apply Mixture Mixture (NeuralNetwork m n o) where
    (>$>) nnp xps =
        let (_,_,_,zps) = feedForward nnp xps
         in zps


--- Backprop ---


{-
--backpropagation :: NeuralNetwork (m ': ms) -> (Mixture :#: m -> Mixture :#: m) -> Differential :#:
backpropagate :: NeuralNetwork (m ': ms) -> Mixture :#: m -> Differential :#: NeuralNetwork (m ': ms)
backpropagate nnp dp =



--- Internal ---


popManifold :: NeuralNetwork (m ': ms) -> (m, NeuralNetwork ms)
popManifold (Layer m ms) = (m,ms)

popNeuralNetwork
    :: (Manifold m, Manifold n, Manifold (NeuralNetwork (n ': ms)))
    => Function Mixture Mixture :#: NeuralNetwork (m ': n ': ms)
    -> (Natural :#: m, NaturalFunction :#: Tensor m n, Function Mixture Mixture :#: NeuralNetwork (n ': ms))
popNeuralNetwork nnp =
    let (m,nn') = popManifold $ manifold nnp
        (n,_) = popManifold nn'
        tns = Tensor m n
        css = coordinates nnp
        (mcs,css') = C.splitAt (dimension m) css
        (mtxcs,nncs') = C.splitAt (dimension tns) css'
        mp = fromCoordinates m mcs
        mtx = fromCoordinates tns mtxcs
        nnp' = fromCoordinates nn' nncs'
     in (mp,mtx,nnp')

feedForward
    :: Function Mixture Mixture :#: NeuralNetwork ms
    -> [Mixture :#: Domain (NeuralNetwork ms)]
    -> [Mixture :#: Responses ms]
feedForward nnp0 xps0 =
    recurse nnp0 xps0 [ chart Mixture . fromCoordinates (Responses $ Layer (manifold xp) Nub) | xp <- xps ]
        where recurse nnp xps rss =
                  let (b,mtx,nnp') = popNeuralNetwork nnp
                      yps = nnp' >$> xps
                   in map (potentialMapping . (<+> b)) $ mtx >$> ys


feedBackward
    :: [Mixture :#: Codomain (NeuralNetwork ms)]
    -> [Mixture :#: Responses ms]
    -> Differential :#: Tangent (Function Mixture Mixture) (NeuralNetwork ms)
feedBackward = undefined

--- Instances ---


-- Responses --

instance Eq (Responses '[]) where
    (==) _ _ = True

instance (Eq m, Eq (NeuralNetwork ms)) => Eq (Responses (m ': ms)) where
    (==) (Responses (Layer m ms)) (Responses (Layer m' ms'))
        | m == m' = ms == ms'
        | otherwise = False

instance Manifold (Responses '[]) where
    dimension _ = 0


instance (Manifold m, Manifold (NeuralNetwork ms)) => Manifold (Responses (m ': ms)) where
    dimension (Responses (Layer m ms)) =  dimension m + dimension ms


-- NeuralNetwork --

instance Eq (NeuralNetwork '[]) where
    (==) _ _ = True

instance (Eq m, Eq (NeuralNetwork ms)) => Eq (NeuralNetwork (m ': ms)) where
    (==) (Layer m ms) (Layer m' ms')
        | m == m' = ms == ms'
        | otherwise = False

instance Manifold (NeuralNetwork '[]) where
    dimension _ = 0

instance Manifold m => Manifold (NeuralNetwork '[m]) where
    dimension _ = 0

instance (Manifold m, Manifold n, Manifold (NeuralNetwork (n ': ms))) => Manifold (NeuralNetwork (m ': n ': ms)) where
    dimension (Layer m (Layer n ms)) =  dimension m + dimension m * dimension n + dimension (Layer n ms)

instance Manifold m => Map (NeuralNetwork '[m]) where
    type Domain (NeuralNetwork '[m]) = m
    domain (Layer m _) = m
    type Codomain (NeuralNetwork '[m]) = m
    codomain (Layer m _) = m

instance (ExponentialFamily m, Manifold n) => Apply Mixture Mixture (NeuralNetwork '[m,n]) where
    (>$>) p xs =
        let (b,mtx,_) = popNeuralNetwork p
         in map (potentialMapping . (<+> b)) $ mtx >$> xs

instance (ExponentialFamily m, Manifold n, Map (NeuralNetwork (n ': ms)))
    => Map (NeuralNetwork (m ': n ': ms)) where
    type Domain (NeuralNetwork (m ': n ': ms)) = Domain (NeuralNetwork (n ': ms))
    domain (Layer _ nn) = domain nn
    type Codomain (NeuralNetwork (m ': n ': ms)) = m
    codomain (Layer m _) = m

instance (ExponentialFamily m, Manifold n, Apply Mixture Mixture (NeuralNetwork (n ': o ': ms)))
    => Apply Mixture Mixture (NeuralNetwork (m ': n ': o ': ms)) where
    (>$>) p xs =
        let (b,mtx,p') = popNeuralNetwork p
            ys = p' >$> xs
         in map (potentialMapping . (<+> b)) $ mtx >$> ys
    -}