content: {-# LANGUAGE AllowAmbiguousTypes #-}
  {-# LANGUAGE BangPatterns        #-}
  {-# LANGUAGE DataKinds           #-}
  {-# LANGUAGE FlexibleContexts    #-}
  {-# LANGUAGE GADTs               #-}
  {-# LANGUAGE KindSignatures      #-}
  {-# LANGUAGE LambdaCase          #-}
  {-# LANGUAGE RankNTypes          #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TypeApplications    #-}
  {-# LANGUAGE TypeInType          #-}
  {-# LANGUAGE TypeOperators       #-}
  
  module TensorOps.Learn.NeuralNet.Recurrent
    ( Network
    , buildNet
    , netParams
    , runNetwork
    , runNetworkSt
    , genNet
    , fullyConnected
    , ffLayer
    , stateless
    , (~*~)
    , (*~)
    , (~*)
    , nmap
    , trainNetwork
    , trainNetwork'
    , networkGradient
    ) where
  
  import           Control.Category
  import           Control.DeepSeq
  import           Control.Monad.Primitive
  import           Control.Monad.State
  import           Data.Kind
  import           Data.Singletons
  import           Data.Singletons.Prelude               (Sing(..))
  import           Data.Type.Combinator
  import           Data.Type.Conjunction
  import           Data.Type.Length                      as TCL
  import           Data.Type.Length.Util                 as TCL
  import           Data.Type.Nat
  import           Data.Type.Product                     as TCP
  import           Data.Type.Product.Util                as TCP
  import           Data.Type.Sing
  import           Data.Type.Uniform
  import           Data.Type.Vector                      as TCV
  import           Data.Type.Vector.Util                 as TCV
  import           Prelude hiding                        ((.), id)
  import           Statistics.Distribution.Normal
  import           System.Random.MWC
  import           TensorOps.Learn.NeuralNet
  import           TensorOps.NatKind
  import           TensorOps.TOp                         as TO
  import           TensorOps.Tensor                      as TT
  import           TensorOps.Types
  import           Type.Class.Higher
  import           Type.Class.Higher.Util
  import           Type.Class.Known
  import           Type.Class.Witness
  import           Type.Family.List
  import           Type.Family.List.Util
  import           Type.Family.Nat
  import           Unsafe.Coerce
  import qualified TensorOps.Learn.NeuralNet.FeedForward as FF
  
  data Network :: ([k] -> Type) -> k -> k -> Type where
      N :: { _nsSs    :: !(Sing ss)
           , _nsPs    :: !(Sing ps)
           , _nOp     :: !(TOp ('[i] ': ss ++ ps) ('[o] ': ss))
           , _nState  :: !(Prod t ss)
           , _nParams :: !(Prod t ps)
           } -> Network t i o
  
  instance NFData1 t => NFData (Network t i o) where
      rnf = \case
        N _ _ o p s -> o `seq` p `deepseq1` s `deepseq1` ()
      {-# INLINE rnf #-}
  
  netParams
      :: Network t i o
      -> (forall ss ps. (SingI ss, SingI ps) => Prod t ss -> Prod t ps -> r)
      -> r
  netParams = \case
      N sS sP _ s p -> \f -> f s p \\ sS \\ sP
  
  buildNet
      :: forall ss ps i o t. (SingI ss, SingI ps)
      => TOp ('[i] ': ss ++ ps) ('[o] ': ss)  -- ^ Tensor operation for network
      -> Prod t ss        -- ^ Initial states
      -> Prod t ps        -- ^ Network parameters
      -> Network t i o
  buildNet = N sing sing
  
  fullyConnected
      :: forall k (i :: k) (o :: k) (m :: Type -> Type) (t :: [k] -> Type).
       ( SingI i
       , SingI o
       , PrimMonad m
       , Tensor t
       )
      => Activation k         -- ^ Activation function for internal state
      -> Gen (PrimState m)
      -> m (Network t i o)
  fullyConnected a g = (\s w w' b -> buildNet @'[ '[o] ] @'[ '[o,o], '[o,i], '[o] ]
                                       fc (s :< Ø) (w' :< w :< b :< Ø)
                       )
            <$> genRand (normalDistr 0 0.5) g
            <*> genRand (normalDistr 0 0.5) g
            <*> genRand (normalDistr 0 0.5) g
            <*> genRand (normalDistr 0 0.5) g
    where
      fc  :: TOp '[ '[i], '[o], [o,o], '[o,i], '[o]] '[ '[o], '[o] ]
      fc = secondOp @'[ '[i] ] (
             firstOp @'[ '[o,i], '[o] ] (TO.swap >>> TO.matVec)
         >>> firstOp @'[ '[o] ]         TO.swap
           )
       >>> firstOp @'[ '[o], '[o] ] (TO.swap >>> TO.matVec)
       >>> TO.add3
       >>> TO.duplicate
       >>> secondOp @'[ '[o] ] (getAct a)
      {-# INLINE fc #-}
  {-# INLINE fullyConnected #-}
  
  -- | Convert a neural network from "TensorOps.Learn.NeuralNet.FeedForward"
  -- to a stateless 'Network'.
  --
  -- Can be thought a functor from the category of stateless neural networks
  -- to the category of stateful (recurrent) ones.
  stateless
      :: FF.Network t i o
      -> Network t i o
  stateless = \case
      FF.N sP o p -> N SNil sP o Ø p
  {-# INLINE stateless #-}
  
  ffLayer
      :: forall i o m t. (SingI i, SingI o, PrimMonad m, Tensor t)
      => Gen (PrimState m)
      -> m (Network t i o)
  ffLayer g = stateless <$> FF.ffLayer g
  {-# INLINE ffLayer #-}
  
  genNet
      :: forall k o i m (t :: [k] -> Type). (SingI o, SingI i, PrimMonad m, Tensor t)
      => [(Integer, (Activation k, Maybe (Activation k)))]
      -> Activation k
      -> Maybe (Activation k)
      -> Gen (PrimState m)
      -> m (Network t i o)
  genNet xs0 f fS g = go sing xs0
    where
      go  :: forall (j :: k). ()
          => Sing j
          -> [(Integer, (Activation k, Maybe (Activation k)))]
          -> m (Network t j o)
      go sj = (\\ sj) $ \case
        []        -> fmap (*~ getAct f) $ case fS of
          Just fS' -> fullyConnected fS' g
          Nothing  -> ffLayer g
        (x,(f', fS')):xs -> withNatKind x $ \sl -> (\\ sl) $ do
          n <- go sl xs
          l <- case fS' of
            Nothing   -> ffLayer g
            Just fS'' -> fullyConnected fS'' g
          return $ l *~ getAct f' ~*~ n
      {-# INLINE go #-}
  {-# INLINE genNet #-}
  
  instance Category (Network t) where
      id  = N SNil SNil idOp Ø Ø
      (.) = flip (~*~)
  
  (~*~)
      :: forall k (t :: [k] -> Type) a b c. ()
      => Network t a b
      -> Network t b c
      -> Network t a c
  (~*~) = \case
    N (sSs1 :: Sing ss1)
      (sPs1 :: Sing ps1)
      (o1   :: TOp ('[a] ': ss1 ++ ps1) ('[b] ': ss1))
      (s1   :: Prod t ss1)
      (p1   :: Prod t ps1) -> \case
      N (sSs2 :: Sing ss2)
        (sPs2 :: Sing ps2)
        (o2   :: TOp ('[b] ': ss2 ++ ps2) ('[c] ': ss2))
        (s2   :: Prod t ss2)
        (p2   :: Prod t ps2) ->
        let lSs1 :: Length ss1
            lSs1 = singLength sSs1
            lSs2 :: Length ss2
            lSs2 = singLength sSs2
            lPs1 :: Length ps1
            lPs1 = singLength sPs1
            lPs2 :: Length ps2
            lPs2 = singLength sPs2
            o :: TOp ('[a] ': (ss2 ++ ss1) ++ (ps1 ++ ps2)) ('[c] ': ss2 ++ ss1)
                      -- all these proofs lol
            o     = (\\ (unsafeCoerce Refl :: ((ss2 ++ ss1) ++ ps1) ++ ps2 :~: (ss2 ++ ss1) ++ (ps1 ++ ps2))) $
                    (\\ (unsafeCoerce Refl :: ((ss1 ++ ps1) ++ ss2) ++ ps2 :~: (ss1 ++ ps1) ++ (ss2 ++ ps2))) $
                    (\\ appendAssoc lSs1 lSs2 lPs2                  ) $
                    (\\ appendAssoc lSs2 lSs1 lPs1                  ) $
                    (\\ (lSs2 `TCL.append'` lSs1) `TCL.append'` lPs1) $
                    (\\ (lSs1 `TCL.append'` lPs1) `TCL.append'` lSs2) $
                    (\\ lSs1                                        ) $
                    (\\ lSs2                                        ) $
                    (\\ lSs1 `TCL.append'` lPs1                     ) $
                    (\\ lSs2 `TCL.append'` lPs2                     ) $
                      secondOp @'[ '[a] ]
                        (firstOp @ps2 (TO.swap' lSs2 (lSs1 `TCL.append'` lPs1))
                        )
                  >>> firstOp @(ss2 ++ ps2) o1
                  >>> secondOp @'[ '[b] ] (TO.swap' lSs1 (lSs2 `TCL.append'` lPs2))
                  >>> firstOp @ss1 o2
        in  N (sSs2 %:++ sSs1)
              (sPs1 %:++ sPs2)
              o
              (s2 `TCP.append'` s1)
              (p1 `TCP.append'` p2)
  infixr 4 ~*~
  {-# INLINE (~*~) #-}
  
  runNetwork
      :: (RealFloat (ElemT t), Tensor t)
      => Network t i o
      -> t '[i]
      -> (t '[o], Network t i o)
  runNetwork (N sS sO o s p) x =
          (\case y :< s' -> (y, N sS sO o s' p))
        . runTOp o
        $ x :< (s `TCP.append'` p)
  {-# INLINE runNetwork #-}
  
  runNetworkSt
      :: (RealFloat (ElemT t), Tensor t, MonadState (Network t i o) m)
      => t '[i]
      -> m (t '[o])
  runNetworkSt x = state $ flip runNetwork x
  
  (~*) :: TOp '[ '[a] ] '[ '[b] ]
       -> Network t b c
       -> Network t a c
  f ~* N sS sO o s p = N sS sO (f *>> o) s p
  infixr 4 ~*
  {-# INLINE (~*) #-}
  
  (*~) :: Network t a b
       -> TOp '[ '[b] ] '[ '[c] ]
       -> Network t a c
  N sS sO o s p *~ f = N sS sO (o >>> firstOp f) s p
  infixl 5 *~
  {-# INLINE (*~) #-}
  
  nmap
       :: SingI o
       => (forall a. RealFloat a => a -> a)
       -> Network t i o
       -> Network t i o
  nmap f n = n *~ TO.map f
  {-# INLINE nmap #-}
  
  netGrad
      :: forall k n (t :: [k] -> Type) (i :: k) (o :: k) ss ps.
       ( Tensor t
       , RealFloat (ElemT t)
       , SingI i
       , SingI o
       )
      => TOp '[ '[o], '[o] ] '[ '[] ]
      -> Vec n (t '[i])
      -> Vec n (t '[o])
      -> Sing ss
      -> Sing ps
      -> TOp ('[i] ': ss ++ ps) ('[o] ': ss)
      -> Prod t ss
      -> Prod t ps
      -> (Vec n (t '[i]), (Prod t ss, Prod t ps))
  netGrad loss xs ys sS sP o s p =
        (prodToVec' I n grI, splitProd @ps lS grSP)
    where
      n :: Nat n
      n = known   \\ xs
      lS :: Length ss
      lS = singLength sS
      lP :: Length ps
      lP = singLength sP
      sO :: Sing (Replicate n '[o])
      sO = replicateSing (sing :: Sing '[o]) n
      lI :: Length (Replicate n '[i])
      lI = replicateLength @'[i] n
      lO :: Length (Replicate n '[o])
      lO = replicateLength @'[o] n
      unrolled
          :: TOp (Replicate n '[i] ++ ss ++ ps) (Replicate n '[o])
      unrolled = (\\ sS %:++ sO) $
                 (\\ sS %:++ sP) $
                   unroll sS sP o n
               >>> TO.drop @(Replicate n '[o]) lS
      o'  :: TOp (Replicate n '[i] ++ ss ++ ps ++ Replicate n '[o]) '[ '[] ]
      o' = (\\ appendAssoc lI (lS `TCL.append'` lP) lO) $
           (\\ appendAssoc lS lP lO) $
           (\\ lI `TCL.append'` lS `TCL.append'` lP) $
           (\\ lO) $
              firstOp @(Replicate n '[o]) unrolled
          >>> rollup @k lS lP loss n
      xs' :: Prod t (Replicate n '[i] ++ ss ++ ps ++ Replicate n '[o])
      xs' = vecToProd' getI (TCV.reverse' xs) `TCP.append'` (s
              `TCP.append'` (
                p `TCP.append'` vecToProd' getI ys
              )
            )
      grad :: Prod t (Replicate n '[i] ++ ss ++ ps)
      grad = (\\ (unsafeCoerce Refl :: Replicate n '[i] ++ ss ++ ps ++ Replicate n '[o]
                                   :~: (Replicate n '[i] ++ ss ++ ps) ++ Replicate n '[o]
                 )
             ) $
        takeProd @(Replicate n '[o]) (lI `TCL.append'` lS `TCL.append'` lP) $ gradTOp o' xs'
      grI   :: Prod t (Replicate n '[i])
      grSP  :: Prod t (ss ++ ps)
      (grI, grSP) = splitProd @(ss ++ ps) lI grad
  {-# INLINE netGrad #-}
  
  trainNetwork'
      :: forall k n (t :: [k] -> Type) (i :: k) (o :: k).
       ( Tensor t
       , RealFloat (ElemT t)
       , SingI i
       , SingI o
       )
      => TOp '[ '[o], '[o] ] '[ '[] ]     -- ^ loss function (input, target)
      -> ElemT t          -- ^ train rate for initial state
      -> ElemT t          -- ^ train rate for network parameters
      -> Vec n (t '[i])   -- ^ inputs
      -> Vec n (t '[o])   -- ^ targets
      -> Network t i o
      -> Network t i o
  trainNetwork' loss rS rP xs ys = \case
      N sS sP o s p ->
        let (gS, gP) = snd $ netGrad loss xs ys sS sP o s p
            s' = map1 (f rS) $ zipProd3 (singProd sS) s gS
            p' = map1 (f rP) $ zipProd3 (singProd sP) p gP
        in  N sS sP o s' p'
    where
      f   :: forall ns. ()
          => ElemT t
          -> (Sing :&: t :&: t) ns
          -> t ns
      f r !(s1 :&: p1 :&: g1) =
        TT.zip (\(!p2) (!g2) -> p2 - r * g2) p1 g1 \\ s1
      {-# INLINE f #-}
  {-# INLINE trainNetwork #-}
  
  trainNetwork
      :: forall k (t :: [k] -> Type) (i :: k) (o :: k).
       ( Tensor t
       , RealFloat (ElemT t)
       , SingI i
       , SingI o
       )
      => TOp '[ '[o], '[o] ] '[ '[] ]     -- ^ loss function (input, target)
      -> ElemT t          -- ^ train rate for initial state
      -> ElemT t          -- ^ train rate for network parameters
      -> [(t '[i], t '[o])]   -- ^ inputs and targets
      -> Network t i o
      -> Network t i o
  trainNetwork loss rS rP xsys n = withV xsys $ \v ->
      let (xs, ys) = TCV.unzip' v
      in  trainNetwork' loss rS rP xs ys n
  
  networkGradient
      :: forall k n (t :: [k] -> Type) (i :: k) (o :: k) r.
       ( Tensor t
       , RealFloat (ElemT t)
       , SingI i
       , SingI o
       )
      => TOp '[ '[o], '[o] ] '[ '[] ]     -- ^ loss function (output, target)
      -> Vec n (t '[i])       -- ^ inputs
      -> Vec n (t '[o])       -- ^ targets
      -> Network t i o
      -> (forall ss ps. (SingI ss, SingI ps) => Vec n (t '[i]) -> Prod t ss -> Prod t ps -> r)
      -> r
  networkGradient loss xs ys = \case
      N sS sP o s p -> \f -> case netGrad loss xs ys sS sP o s p of
        (gI, (gS, gP)) -> f gI gS gP \\ sS \\ sP
  {-# INLINE networkGradient #-}
  
  
  unroll
      :: forall ss ps i o n. (SingI (ss ++ ps), SingI i)
      => Sing ss
      -> Sing ps
      -> TOp ('[i] ': ss ++ ps) ('[o] ': ss)
      -> Nat n
      -> TOp (Replicate n '[i] ++ ss ++ ps) (ss ++ Replicate n '[o])
  unroll sS sP o = \case
      Z_              -> TO.take lS lP    \\ appendNil lS
      S_ (m :: Nat m) -> (\\ (unsafeCoerce Refl :: Replicate m '[i] ++ '[i] ': ss ++ ps
                                               :~: '[i] ': Replicate m '[i] ++ ss ++ ps
                             )
                         ) $
                         (\\ (unsafeCoerce Refl :: (Replicate m '[i] ++ ss ++ ps) ++ '[ '[o] ]
                                               :~: Replicate m '[i] ++ ss ++ (ps >: '[o])
                             )
                         ) $
                         (\\ (unsafeCoerce Refl :: (ss ++ Replicate m '[o]) ++ '[ '[o] ]
                                               :~: ss ++ '[o] ': Replicate m '[o]
                             )
                         ) $
                         (\\ appendAssoc lS lP (LS LZ :: Length '[ '[o] ])) $
                         (\\ appendSnoc lP (Proxy @'[o])) $
                         (\\ replicateLength @'[i] m `TCL.append'` lS `TCL.append'` lP) $
                         (\\ lS `TCL.append'` replicateLength @'[o] m) $
                         (\\ replicateLength @'[i] m) $
                         (\\ lS) $
            secondOp @(Replicate m '[i]) @('[i] ': ss ++ ps) @(ss ++ ps >: '[o]) (
                  (o &&& TO.drop @ps @('[i] ': ss) (LS lS))
              >>> TO.swap' (LS LZ) (lS `TCL.append'` lP)
            )
        >>> firstOp @'[ '[o] ] @(Replicate m '[i] ++ ss ++ ps) @(ss ++ Replicate m '[o]) (
              unroll sS sP o m
            )
    where
      lS :: Length ss
      lS = singLength sS
      lP :: Length ps
      lP = singLength sP
  {-# INLINE unroll #-}
  
  
  rollup
      :: forall k (ss :: [[k]]) (ps :: [[k]]) (o :: k) (n :: N). ()
      => Length ss
      -> Length ps
      -> TOp '[ '[o], '[o] ] '[ '[] ]
      -> Nat n
      -> TOp (Replicate n '[o] ++ Replicate n '[o]) '[ '[] ]
  rollup lS lP loss = \case
      Z_                     -> TO.konst (US UØ) 0
      S_ Z_                  -> loss
      S_ (m@(S_ _) :: Nat m) ->
        let lO :: Length (Replicate m '[o])
            lO = replicateLength @'[o] m
        in  (\\ (unsafeCoerce Refl :: Replicate m '[o] ++ '[o] ': '[o] ': Replicate m '[o]
                                  :~: '[o] ': Replicate m '[o] ++ '[o] ': Replicate m '[o]
                )
              ) $
              (\\ appendSnoc lO (Proxy @'[])) $
              (\\ lO) $
              (\\ lO `TCL.append'` lO) $
              (\\ appendAssoc lO lO (LS LZ :: Length '[ '[] ])) $
                secondOp @(Replicate m '[o]) @('[o] ': '[o] ': Replicate m '[o]) @(Replicate m '[o] >: '[]) (
                    firstOp @(Replicate m '[o]) @('[ '[o], '[o] ]) @('[ '[] ]) loss
                  >>> TO.swap' @'[ '[] ] @(Replicate m '[o]) (LS LZ) lO
                )
              >>> firstOp @'[ '[] ] @(Replicate m '[o] ++ Replicate m '[o]) @'[ '[] ] (
                    rollup lS lP loss m
                  )
              >>> TO.add
  {-# INLINE rollup #-}
  