content: {-# LANGUAGE CPP                   #-}
  {-# LANGUAGE DataKinds             #-}
  {-# LANGUAGE GADTs                 #-}
  {-# LANGUAGE TypeOperators         #-}
  {-# LANGUAGE TypeFamilies          #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  {-# LANGUAGE FlexibleContexts      #-}
  {-# LANGUAGE FlexibleInstances     #-}
  {-# LANGUAGE EmptyDataDecls        #-}
  {-# LANGUAGE RankNTypes            #-}
  {-# LANGUAGE BangPatterns          #-}
  {-# LANGUAGE ScopedTypeVariables   #-}
  {-# LANGUAGE UndecidableInstances  #-}
  
  module Grenade.Recurrent.Core.Network (
      Recurrent
    , FeedForward
  
    , RecurrentNetwork (..)
    , RecurrentInputs (..)
    , RecurrentTape (..)
    , RecurrentGradient (..)
  
    , randomRecurrent
    , runRecurrent
    , runRecurrent'
    , applyRecurrentUpdate
    ) where
  
  
  import           Control.Monad.Random ( MonadRandom )
  import           Data.Serialize
  
  import           Data.Kind (Type)
  
  import           Grenade.Core
  import           Grenade.Recurrent.Core.Layer
  import           Prelude.Singletons
  
  -- | Witness type to say indicate we're building up with a normal feed
  --   forward layer.
  data FeedForward :: Type -> Type
  -- | Witness type to say indicate we're building up with a recurrent layer.
  data Recurrent :: Type -> Type
  
  -- | Type of a recurrent neural network.
  --
  --   The [Type] type specifies the types of the layers.
  --
  --   The [Shape] type specifies the shapes of data passed between the layers.
  --
  --   The definition is similar to a Network, but every layer in the
  --   type is tagged by whether it's a FeedForward Layer of a Recurrent layer.
  --
  --   Often, to make the definitions more concise, one will use a type alias
  --   for these empty data types.
  data RecurrentNetwork :: [Type] -> [Shape] -> Type where
    RNil   :: SingI i
           => RecurrentNetwork '[] '[i]
  
    (:~~>) :: (SingI i, Layer x i h)
           => !x
           -> !(RecurrentNetwork xs (h ': hs))
           -> RecurrentNetwork (FeedForward x ': xs) (i ': h ': hs)
  
    (:~@>) :: (SingI i, RecurrentLayer x i h)
           => !x
           -> !(RecurrentNetwork xs (h ': hs))
           -> RecurrentNetwork (Recurrent x ': xs) (i ': h ': hs)
  infixr 5 :~~>
  infixr 5 :~@>
  
  -- | Gradient of a network.
  --
  --   Parameterised on the layers of the network.
  data RecurrentGradient :: [Type] -> Type where
     RGNil  :: RecurrentGradient '[]
  
     (://>) :: UpdateLayer x
            => Gradient x
            -> RecurrentGradient xs
            -> RecurrentGradient (phantom x ': xs)
  
  -- | Recurrent inputs (sideways shapes on an imaginary unrolled graph)
  --   Parameterised on the layers of a Network.
  data RecurrentInputs :: [Type] -> Type where
     RINil   :: RecurrentInputs '[]
  
     (:~~+>) :: (UpdateLayer x, Fractional (RecurrentInputs xs))
             => ()                      -> !(RecurrentInputs xs) -> RecurrentInputs (FeedForward x ': xs)
  
     (:~@+>) :: (Fractional (RecurrentShape x), Fractional (RecurrentInputs xs), RecurrentUpdateLayer x)
             => !(RecurrentShape x) -> !(RecurrentInputs xs) -> RecurrentInputs (Recurrent x ': xs)
  
  -- | All the information required to backpropogate
  --   through time safely.
  --
  --   We index on the time step length as well, to ensure
  --   that that all Tape lengths are the same.
  data RecurrentTape :: [Type] -> [Shape] -> Type where
     TRNil  :: SingI i
            => RecurrentTape '[] '[i]
  
     (:\~>) :: Tape x i h
            -> !(RecurrentTape xs (h ': hs))
            -> RecurrentTape (FeedForward x ': xs) (i ': h ': hs)
  
     (:\@>) :: RecTape x i h
            -> !(RecurrentTape xs (h ': hs))
            -> RecurrentTape (Recurrent x ': xs) (i ': h ': hs)
  
  
  runRecurrent :: forall shapes layers.
                  RecurrentNetwork layers shapes
               -> RecurrentInputs layers
               -> S (Head shapes)
               -> (RecurrentTape layers shapes, RecurrentInputs layers, S (Last shapes))
  runRecurrent =
    go
      where
    go  :: forall js sublayers. (Last js ~ Last shapes)
        => RecurrentNetwork sublayers js
        -> RecurrentInputs sublayers
        -> S (Head js)
        -> (RecurrentTape sublayers js, RecurrentInputs sublayers, S (Last js))
    go (!layer :~~> n) (() :~~+> nIn) !x
        = let (!tape, !forwards) = runForwards layer x
  
              -- recursively run the rest of the network, and get the gradients from above.
              (!newFN, !ig, !answer) = go n nIn forwards
          in (tape :\~> newFN, () :~~+> ig, answer)
  
    -- This is a recurrent layer, so we need to do a scan, first input to last, providing
    -- the recurrent shape output to the next layer.
    go (layer :~@> n) (recIn :~@+> nIn) !x
        = let (tape, shape, forwards) = runRecurrentForwards layer recIn x
              (newFN, ig, answer) = go n nIn forwards
          in (tape :\@> newFN, shape :~@+> ig, answer)
  
    -- Handle the output layer, bouncing the derivatives back down.
    -- We may not have a target for each example, so when we don't use 0 gradient.
    go RNil RINil !x
      = (TRNil, RINil, x)
  
  runRecurrent' :: forall layers shapes.
                   RecurrentNetwork layers shapes
                -> RecurrentTape layers shapes
                -> RecurrentInputs layers
                -> S (Last shapes)
                -> (RecurrentGradient layers, RecurrentInputs layers, S (Head shapes))
  runRecurrent' net tapes r o =
    go net tapes r
      where
    -- We have to be careful regarding the direction of the lists
    -- Inputs come in forwards, but our return value is backwards
    -- through time.
    go  :: forall js ss. (Last js ~ Last shapes)
        => RecurrentNetwork ss js
        -> RecurrentTape ss js
        -> RecurrentInputs ss
        -> (RecurrentGradient ss, RecurrentInputs ss, S (Head js))
    -- This is a simple non-recurrent layer
    -- Run the rest of the network, update with the tapes and gradients
    go (!layer :~~> n) (!tape :\~> nTapes) (() :~~+> nRecs) =
      let (!gradients, !rins, !feed)  = go n nTapes nRecs
          (!grad, !back)              = runBackwards layer tape feed
      in  (grad ://> gradients, () :~~+> rins, back)
  
    -- This is a recurrent layer
    -- Run the rest of the network, scan over the tapes in reverse
    go (!layer :~@> n) (!tape :\@> nTapes) (!recGrad :~@+> nRecs) =
      let (!gradients, !rins, !feed)    = go n nTapes nRecs
          (!grad, !sidegrad, !back)     = runRecurrentBackwards layer tape recGrad feed
      in  (grad ://> gradients, sidegrad :~@+> rins, back)
  
    -- End of the road, so we reflect the given gradients backwards.
    -- Crucially, we reverse the list, so it's backwards in time as
    -- well.
    go !RNil !TRNil !RINil
      = (RGNil, RINil, o)
  
  -- | Apply a batch of gradients to the network
  --   Uses runUpdates which can be specialised for
  --   a layer.
  applyRecurrentUpdate :: LearningParameters
                       -> RecurrentNetwork layers shapes
                       -> RecurrentGradient layers
                       -> RecurrentNetwork layers shapes
  applyRecurrentUpdate rate (layer :~~> rest) (gradient ://> grest)
    = runUpdate rate layer gradient :~~> applyRecurrentUpdate rate rest grest
  
  applyRecurrentUpdate rate (layer :~@> rest) (gradient ://> grest)
    = runUpdate rate layer gradient :~@> applyRecurrentUpdate rate rest grest
  
  applyRecurrentUpdate _ RNil RGNil
    = RNil
  
  
  instance Show (RecurrentNetwork '[] '[i]) where
    show RNil = "NNil"
  instance (Show x, Show (RecurrentNetwork xs rs)) => Show (RecurrentNetwork (FeedForward x ': xs) (i ': rs)) where
    show (x :~~> xs) = show x ++ "\n~~>\n" ++ show xs
  instance (Show x, Show (RecurrentNetwork xs rs)) => Show (RecurrentNetwork (Recurrent x ': xs) (i ': rs)) where
    show (x :~@> xs) = show x ++ "\n~~>\n" ++ show xs
  
  
  -- | A network can easily be created by hand with (:~~>) and (:~@>), but an easy way to initialise a random
  --   recurrent network and a set of random inputs for it is with the randomRecurrent.
  class CreatableRecurrent (xs :: [Type]) (ss :: [Shape]) where
    -- | Create a network of the types requested
    randomRecurrent :: MonadRandom m => m (RecurrentNetwork xs ss)
  
  instance SingI i => CreatableRecurrent '[] '[i] where
    randomRecurrent =
      return RNil
  
  instance (SingI i, Layer x i o, CreatableRecurrent xs (o ': rs)) => CreatableRecurrent (FeedForward x ': xs) (i ': o ': rs) where
    randomRecurrent = do
      thisLayer     <- createRandom
      rest          <- randomRecurrent
      return (thisLayer :~~> rest)
  
  instance (SingI i, RecurrentLayer x i o, CreatableRecurrent xs (o ':  rs)) => CreatableRecurrent (Recurrent x ': xs) (i ': o ': rs) where
    randomRecurrent = do
      thisLayer     <- createRandom
      rest          <- randomRecurrent
      return (thisLayer :~@> rest)
  
  -- | Add very simple serialisation to the recurrent network
  instance SingI i => Serialize (RecurrentNetwork '[] '[i]) where
    put RNil = pure ()
    get = pure RNil
  
  instance (SingI i, Layer x i o, Serialize x, Serialize (RecurrentNetwork xs (o ': rs))) => Serialize (RecurrentNetwork (FeedForward x ': xs) (i ': o ': rs)) where
    put (x :~~> r) = put x >> put r
    get = (:~~>) <$> get <*> get
  
  instance (SingI i, RecurrentLayer x i o, Serialize x, Serialize (RecurrentNetwork xs (o ': rs))) => Serialize (RecurrentNetwork (Recurrent x ': xs) (i ': o ': rs)) where
    put (x :~@> r) = put x >> put r
    get = (:~@>) <$> get <*> get
  
  instance (Serialize (RecurrentInputs '[])) where
    put _ = return ()
    get = return RINil
  
  instance (UpdateLayer x, Serialize (RecurrentInputs ys), Fractional (RecurrentInputs ys)) => (Serialize (RecurrentInputs (FeedForward x ': ys))) where
    put ( () :~~+> rest) = put rest
    get = ( () :~~+> ) <$> get
  
  instance (Serialize (RecurrentShape x), Fractional (RecurrentShape x), RecurrentUpdateLayer x, Serialize (RecurrentInputs ys), Fractional (RecurrentInputs ys)) => (Serialize (RecurrentInputs (Recurrent x ': ys))) where
    put ( i :~@+> rest ) = put i >> put rest
    get = (:~@+>) <$> get <*> get
  
  
  -- Num instance for `RecurrentInputs layers`
  -- Not sure if this is really needed, as I only need a `fromInteger 0` at
  -- the moment for training, to create a null gradient on the recurrent
  -- edge.
  --
  -- It does raise an interesting question though? Is a 0 gradient actually
  -- the best?
  --
  -- I could imaging that weakly push back towards the optimum input could
  -- help make a more stable generator.
  instance (Num (RecurrentInputs '[])) where
    (+) _ _  = RINil
    (-) _ _  = RINil
    (*) _ _  = RINil
    abs _    = RINil
    signum _ = RINil
    fromInteger _ = RINil
  
  instance (UpdateLayer x, Fractional (RecurrentInputs ys)) => (Num (RecurrentInputs (FeedForward x ': ys))) where
    (+) (() :~~+> x) (() :~~+> y)  = () :~~+> (x + y)
    (-) (() :~~+> x) (() :~~+> y)  = () :~~+> (x - y)
    (*) (() :~~+> x) (() :~~+> y)  = () :~~+> (x * y)
    abs (() :~~+> x)      = () :~~+> abs x
    signum (() :~~+> x)   = () :~~+> signum x
    fromInteger x         = () :~~+> fromInteger x
  
  instance (Fractional (RecurrentShape x), RecurrentUpdateLayer x, Fractional (RecurrentInputs ys)) => (Num (RecurrentInputs (Recurrent x ': ys))) where
    (+) (x :~@+> x') (y :~@+> y')  = (x + y) :~@+> (x' + y')
    (-) (x :~@+> x') (y :~@+> y')  = (x - y) :~@+> (x' - y')
    (*) (x :~@+> x') (y :~@+> y')  = (x * y) :~@+> (x' * y')
    abs (x :~@+> x')      = abs x :~@+> abs x'
    signum (x :~@+> x')   = signum x :~@+> signum x'
    fromInteger x         = fromInteger x :~@+> fromInteger x
  
  instance (Fractional (RecurrentInputs '[])) where
    (/) _ _        = RINil
    recip _        = RINil
    fromRational _ = RINil
  
  instance (UpdateLayer x, Fractional (RecurrentInputs ys)) => (Fractional (RecurrentInputs (FeedForward x ': ys))) where
    (/) (() :~~+> x) (() :~~+> y)  = () :~~+> (x / y)
    recip (() :~~+> x)    = () :~~+> recip x
    fromRational x        = () :~~+> fromRational x
  
  instance (Fractional (RecurrentShape x), RecurrentUpdateLayer x, Fractional (RecurrentInputs ys)) => (Fractional (RecurrentInputs (Recurrent x ': ys))) where
    (/) (x :~@+> x') (y :~@+> y') = (x / y) :~@+> (x' / y')
    recip (x :~@+> x')    = recip x :~@+> recip x'
    fromRational x        = fromRational x :~@+> fromRational x
  
  
  -- | Ultimate composition.
  --
  --   This allows a complete network to be treated as a layer in a larger network.
  instance CreatableRecurrent sublayers subshapes => UpdateLayer (RecurrentNetwork sublayers subshapes) where
    type Gradient (RecurrentNetwork sublayers subshapes) = RecurrentGradient sublayers
    runUpdate    = applyRecurrentUpdate
    createRandom = randomRecurrent
  
  -- | Ultimate composition.
  --
  --   This allows a complete network to be treated as a layer in a larger network.
  instance CreatableRecurrent sublayers subshapes => RecurrentUpdateLayer (RecurrentNetwork sublayers subshapes) where
    type RecurrentShape (RecurrentNetwork sublayers subshapes) = RecurrentInputs sublayers
  
  -- | Ultimate composition.
  --
  --   This allows a complete network to be treated as a layer in a larger network.
  instance ( CreatableRecurrent sublayers subshapes
           , i ~ (Head subshapes), o ~ (Last subshapes)
           , Num (RecurrentShape (RecurrentNetwork sublayers subshapes))
           ) => RecurrentLayer (RecurrentNetwork sublayers subshapes) i o where
    type RecTape (RecurrentNetwork sublayers subshapes) i o = RecurrentTape sublayers subshapes
    runRecurrentForwards = runRecurrent
    runRecurrentBackwards = runRecurrent'
  