content: {- | This module implements reverse-mode automatic differentiation.
  
  Machine learning and training of models are based on calculating gradients of operations.
  This can be done symbolically by dynamically creating a graph of all operations,
  which is then traversed to obtain the gradient.
  
  "Synapse" provides several operations that support automatic differentiation,
  but you could easily extend list of those: you just need to define function
  that returns 'Symbol' with correct local gradients.
  You can check out implementations in the source to give yourself a reference
  and read more about it in 'Symbol' datatype docs.
  -}
  
  
  -- 'FlexibleInstances' and 'TypeFamilies' are needed to instantiate 'Indexable', 'ElementwiseScalarOps', 'SingletonOps', 'VecOps', 'MatOps' typeclasses.
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE TypeFamilies      #-}
  
  
  module Synapse.Autograd
      ( -- * 'Symbolic' and 'Symbol'
  
        Symbolic (symbolicZero, symbolicOne, symbolicN)
  
      , SymbolIdentifier (SymbolIdentifier, unSymbolIdentifier)
  
      , Symbol (Symbol, symbolIdentifier, unSymbol, symbolGradients)
      , SymbolVec
      , SymbolMat
  
      , symbol
      , constSymbol
      , renameSymbol
  
      , symbolicUnaryOp
      , symbolicBinaryOp
  
        -- * 'Gradients' calculation
  
      , Gradients (unGradients)
      , allGradients
      , getGradientsOf
      , wrt
      , nthPartialGradient
      , nthGradient
      ) where
  
  
  import Synapse.Tensors (DType, Indexable(..), ElementwiseScalarOps(..), ElementwiseScalarOps(..), SingletonOps(..), VecOps(..), MatOps(..))
  
  import Synapse.Tensors.Vec (Vec)
  import qualified Synapse.Tensors.Vec as V
  
  import Synapse.Tensors.Mat (Mat)
  import qualified Synapse.Tensors.Mat as M
  
  import Data.Foldable (foldl')
  import Data.String (IsString(..))
  
  import Data.Hashable (Hashable(..))
  
  import qualified Data.HashMap.Lazy as HM
  
  
  {- | 'Symbolic' typeclass describes types with few properties that are needed for autogradient.
  
  Members of this typeclass could have default implementation due to 'Num', but such implementation is not always correct.
  'Synapse.Tensors.Vec.Vec's and 'Synapse.Tensors.Mat.Mat's do not have only one zero or identity element, and so numerical literal is not enough.
  'symbolicZero' and 'symbolicOne' function additionally take reference value to consider dimensions.
  Absence of default implementations forces to manually ensure correctness of those functions.
  
  "Synapse" provides implementations for primitive types ('Int', 'Float', 'Double'),
  and for containers types ('Synapse.Tensors.Vec.Vec', 'Synapse.Tensors.Vec.Vec').
  
  Detailed laws of 'Symbolic' properties are in the docs for associated functions.
  -}
  class (Eq a, Num a) => Symbolic a where
      -- | Returns additive and multiplicative (elementwise) zero element. Argument is passed for the reference of the dimension.
      symbolicZero :: a -> a
  
      -- | Returns multiplicative (elementwise) identity element. Argument is passed for the reference of the dimension.
      symbolicOne :: a -> a
  
      -- | Returns what could be considered @N@ constant (sum of @N@ 'symbolicOne's). Argument is passed for the reference of the dimension.
      symbolicN :: Int -> a -> a
      symbolicN n c
          | n < 0     = negate $ go (abs n) c
          | n == 0    = symbolicZero c
          | otherwise = go n c
        where
          go 0 x = x
          go i x = go (i - 1) (x + symbolicOne x)
  
  -- Instances
  
  instance Symbolic Int where
      symbolicZero = const 0
      symbolicOne = const 1
      symbolicN = const
  
  instance Symbolic Float where
      symbolicZero = const 0.0
      symbolicOne = const 1.0
      symbolicN n = const (fromIntegral n)
  
  instance Symbolic Double where
      symbolicZero = const 0.0
      symbolicOne = const 1.0
      symbolicN n = const (fromIntegral n)
  
  
  instance Symbolic a => Symbolic (Vec a) where
      symbolicZero reference = V.replicate (V.size reference) 0
      symbolicOne reference = V.replicate (V.size reference) 1
      symbolicN n reference = V.replicate (V.size reference) (fromIntegral n)
  
  instance Symbolic a => Symbolic (Mat a) where
      symbolicZero reference = M.replicate (M.size reference) 0
      symbolicOne reference = M.replicate (M.size reference) 1
      symbolicN n reference = M.replicate (M.size reference) (fromIntegral n)
  
  
  -- | 'SymbolIdentifier' is a newtype that wraps string, which needs to uniquely represent symbol.
  newtype SymbolIdentifier = SymbolIdentifier 
      { unSymbolIdentifier :: String  -- ^ Identifier of a symbol.
      } deriving (Eq, Show)
  
  instance IsString SymbolIdentifier where
      fromString = SymbolIdentifier
  
  instance Semigroup SymbolIdentifier where
      (<>) (SymbolIdentifier a) (SymbolIdentifier b) = SymbolIdentifier $ a <> b
  
  instance Monoid SymbolIdentifier where
      mempty = SymbolIdentifier ""
  
  {- | Datatype that represents symbol variable (variable which operations are recorded to symbolically obtain derivatives).
  
  Any operation returning @Symbol a@ where @a@ is 'Symbolic' could be autogradiented - returned 'Symbol' has 'symbolGradients' list,
  which allows "Synapse" to build a graph of computation and obtain needed gradients.
  'symbolGradients' list contains pairs: first element in that pair is symbol wrt which you can take gradient and
  the second element is closure that represents chain rule - it takes incoming local gradient of said symbol and multiplies it by local derivative.
  You can check out implementations of those operations in the source to give yourself a reference.
  -}
  data Symbol a = Symbol
      { symbolIdentifier :: SymbolIdentifier                    -- ^ Name of a symbol (identifier for differentiation).
      , unSymbol         :: a                                   -- ^ Value of a symbol.
      , symbolGradients  :: [(Symbol a, Symbol a -> Symbol a)]  -- ^ List of gradients (wrt to what Symbol and closure to calculate gradient). 
      }
  
  -- | Creates new symbol that refers to a variable (so it must have a name to be able to be differentiated wrt).
  symbol :: SymbolIdentifier -> a -> Symbol a
  symbol name value = Symbol name value []
  
  -- | Creates new symbol that refers to constant (so it does not have name and thus its gradients are not saved).
  constSymbol :: a -> Symbol a
  constSymbol = symbol mempty
  
  -- | Renames symbol which allows differentiating wrt it. Note: renaming practically creates new symbol for the gradient calculation.
  renameSymbol :: SymbolIdentifier -> Symbol a -> Symbol a
  renameSymbol name (Symbol _ value localGradients) = Symbol name value localGradients
  
  
  -- | @SymbolVec a@ type alias stands for @Symbol (Vec a)@.
  type SymbolVec a = Symbol (Vec a)
  
  -- | @SymbolMat a@ type alias stands for @Symbol (Mat a)@.
  type SymbolMat a = Symbol (Mat a)
  
  
  -- Typeclasses
  
  instance Show a => Show (Symbol a) where
      show (Symbol name value _) = "Symbol " ++ show name ++ ": " ++ show value
  
  
  instance Eq (Symbol a) where
      (==) (Symbol name1 _ _) (Symbol name2 _ _) = name1 == name2
  
  
  instance Hashable (Symbol a) where
      hashWithSalt salt (Symbol (SymbolIdentifier name) _ _) = hashWithSalt salt name
  
  
  -- Symbolic symbols
  
  instance Symbolic a => Symbolic (Symbol a) where
      symbolicZero x = constSymbol $ symbolicZero $ unSymbol x
      symbolicOne x = constSymbol $ symbolicOne $ unSymbol x
  
  
  type instance DType (SymbolVec a) = a
  
  type instance DType (SymbolMat a) = a
  
  
  -- | Converts unary operation into symbolic one.
  symbolicUnaryOp :: (a -> a) -> Symbol a -> [(Symbol a, Symbol a -> Symbol a)] -> Symbol a
  symbolicUnaryOp op x = Symbol mempty (op (unSymbol x))
  
  -- | Converts binary operation into symbolic one.
  symbolicBinaryOp :: (a -> a -> a) -> Symbol a -> Symbol a -> [(Symbol a, Symbol a -> Symbol a)] -> Symbol a
  symbolicBinaryOp op a b = Symbol mempty (op (unSymbol a) (unSymbol b))
  
  instance Symbolic a => Num (Symbol a) where
      (+) a b = symbolicBinaryOp (+) a b [(a, id), (b, id)]
      (-) a b = symbolicBinaryOp (-) a b [(a, id), (b, negate)]
      negate x = symbolicUnaryOp negate x [(x, negate)]
      (*) a b = symbolicBinaryOp (*) a b [(a, (* b)), (b, (a *))]
      abs x = symbolicUnaryOp abs x [(x, signum)]
      signum x = symbolicUnaryOp signum x [(x, id)]
      fromInteger = constSymbol . fromInteger
  
  instance (Symbolic a, Fractional a) => Fractional (Symbol a) where
      (/) a b = symbolicBinaryOp (/) a b [(a, (/ b)), (b, (* (negate a / (b * b))))]
      recip x = symbolicUnaryOp recip x [(x, (* (negate (symbolicOne x) / (x * x))))]
      fromRational = constSymbol . fromRational
  
  instance (Symbolic a, Floating a) => Floating (Symbol a) where
      pi = constSymbol pi
      (**) a b = symbolicBinaryOp (**) a b [(a, (* (b * a ** (b - symbolicOne b)))), (b, (* (a ** b * log a)))]
      sqrt x = symbolicUnaryOp sqrt x [(x, (* (recip $ symbolicN 2 x * sqrt x)))]
      exp x = symbolicUnaryOp exp x [(x, (* exp x))]
      log x = symbolicUnaryOp log x [(x, (* recip x))]
      sin x = symbolicUnaryOp sin x [(x, (* cos x))]
      cos x = symbolicUnaryOp cos x [(x, (* negate (sin x)))]
      asin x = symbolicUnaryOp asin x [(x, (* recip (sqrt (symbolicOne x - x * x))))]
      acos x = symbolicUnaryOp acos x [(x, (* negate (recip (sqrt (symbolicOne x - x * x)))))]
      atan x = symbolicUnaryOp atan x [(x, (* recip (symbolicOne x + x * x)))]
      sinh x = symbolicUnaryOp sinh x [(x, (* cosh x))]
      cosh x = symbolicUnaryOp cosh x [(x, (* sinh x))]
      asinh x = symbolicUnaryOp asinh x [(x, (* recip (sqrt (symbolicOne x + x * x))))]
      acosh x = symbolicUnaryOp acosh x [(x, (* recip (sqrt (x * x - symbolicOne x))))]
      atanh x = symbolicUnaryOp atanh x [(x, (* recip (symbolicOne x - x * x)))]
  
  instance Symbolic a => ElementwiseScalarOps (Symbol (Vec a)) where
      (+.) x n = x + constSymbol (V.replicate (V.size $ unSymbol x) n)
      (-.) x n = x - constSymbol (V.replicate (V.size $ unSymbol x) n)
      (*.) x n = x * constSymbol (V.replicate (V.size $ unSymbol x) n)
      (/.) x n = x / constSymbol (V.replicate (V.size $ unSymbol x) n)
      (**.) x n = x ** constSymbol (V.replicate (V.size $ unSymbol x) n)
  
      elementsMin x n = symbolicUnaryOp (`elementsMin` n) x
                        [(x, (* constSymbol (V.generate (V.size $ unSymbol x) $ \i -> if unsafeIndex (unSymbol x) i <= n then 1 else 0)))]
      elementsMax x n = symbolicUnaryOp (`elementsMax` n) x
                        [(x, (* constSymbol (V.generate (V.size $ unSymbol x) $ \i -> if unsafeIndex (unSymbol x) i >= n then 1 else 0)))]
  
  instance Symbolic a => ElementwiseScalarOps (SymbolMat a) where
      (+.) x n = x + constSymbol (M.replicate (M.size $ unSymbol x) n)
      (-.) x n = x - constSymbol (M.replicate (M.size $ unSymbol x) n)
      (*.) x n = x * constSymbol (M.replicate (M.size $ unSymbol x) n)
      (/.) x n = x / constSymbol (M.replicate (M.size $ unSymbol x) n)
      (**.) x n = x ** constSymbol (M.replicate (M.size $ unSymbol x) n)
  
      elementsMin x n = symbolicUnaryOp (`elementsMin` n) x
                        [(x, (* constSymbol (M.generate (M.size $ unSymbol x) $ \i -> if unsafeIndex (unSymbol x) i <= n then 1 else 0)))]
      elementsMax x n = symbolicUnaryOp (`elementsMax` n) x
                        [(x, (* constSymbol (M.generate (M.size $ unSymbol x) $ \i -> if unsafeIndex (unSymbol x) i >= n then 1 else 0)))]
  
  instance Symbolic a => SingletonOps (SymbolVec a) where
      singleton = constSymbol . singleton
      isSingleton = isSingleton . unSymbol
      unSingleton = unSingleton . unSymbol
  
      extendSingleton vec reference = constSymbol $ extendSingleton (unSymbol vec) (unSymbol reference)
  
      elementsSum x = symbolicUnaryOp elementsSum x [(x, (`extendSingleton` x))]
      elementsProduct x = let innerProduct = unSingleton $ elementsProduct $ unSymbol x
                          in symbolicUnaryOp elementsProduct x [(x, \path -> extendSingleton path x * constSymbol (V.generate (V.size $ unSymbol x) $ \i -> innerProduct / unsafeIndex (unSymbol x) i))]
  
      mean x = symbolicUnaryOp mean x [(x, \path -> extendSingleton path x /. fromIntegral (V.size (unSymbol x)))]
  
      norm x = symbolicUnaryOp norm x [(x, \path -> extendSingleton path x * (x /. unSingleton (norm $ unSymbol x)))]
  
  instance Symbolic a => SingletonOps (SymbolMat a) where
      singleton = constSymbol . singleton
      isSingleton = isSingleton . unSymbol
      unSingleton = unSingleton . unSymbol
  
      extendSingleton mat reference = constSymbol $ extendSingleton (unSymbol mat) (unSymbol reference)
  
      elementsSum x = symbolicUnaryOp elementsSum x [(x, (`extendSingleton` x))]
      elementsProduct x = let innerProduct = unSingleton $ elementsProduct $ unSymbol x
                          in symbolicUnaryOp elementsProduct x [(x, \path -> extendSingleton path x * constSymbol (M.generate (M.size $ unSymbol x) $ \i -> innerProduct / unsafeIndex (unSymbol x) i))]
  
      mean x = symbolicUnaryOp mean x [(x, \path -> extendSingleton path x /. fromIntegral (M.nElements $ unSymbol x))]
  
      norm x = symbolicUnaryOp norm x [(x, \path -> extendSingleton path x * (x /. unSingleton (norm $ unSymbol x)))]
  
  instance Symbolic a => VecOps (SymbolVec a) where
      dot a b = elementsSum $ a * b
  
  instance Symbolic a => MatOps (SymbolMat a) where
      transpose x = symbolicUnaryOp M.transpose x [(x, (* transpose x))]
  
      addMatRow mat row = symbolicBinaryOp addMatRow mat row [(mat, id), (row, constSymbol . M.rowVec . flip M.indexRow 0 . unSymbol)]
  
      matMul a b = symbolicBinaryOp M.matMul a b [(a, (`matMul` transpose b)), (b, (transpose a `matMul`))]
  
  
  -- Gradients calculation
  
  -- | 'Gradients' datatype holds all gradients of one symbol with respect to other symbols.
  newtype Gradients a = Gradients
      { unGradients :: HM.HashMap (Symbol a) (Symbol a)  -- ^ Map of gradients.
      }
  
  -- | Returns key-value pairs of all gradients of symbol.
  allGradients :: Gradients a -> [(Symbol a, Symbol a)]
  allGradients = HM.toList . unGradients
  
  
  -- Typeclasses
  
  instance Show a => Show (Gradients a) where
      show gradients = show $ allGradients gradients
  
  
  -- | Generates 'Gradients' for given symbol.
  getGradientsOf :: Symbolic a => Symbol a -> Gradients a
  getGradientsOf differentiatedSymbol = Gradients $ HM.insert differentiatedSymbol wrtItself $
                                                    HM.delete (Symbol mempty undefined []) $
                                                    go HM.empty differentiatedSymbol wrtItself
    where
      wrtItself = symbolicOne differentiatedSymbol
  
      go grads s pathValue =
          foldr (\(child, mulPath) grad -> let pathValue' = mulPath pathValue
                                               grad' = HM.alter (\e -> Just $ case e of
                                                                                  Nothing -> pathValue'
                                                                                  Just x  -> x + pathValue'
                                                                ) child grad
                                           in go grad' child pathValue') grads (symbolGradients s)
  
  -- | Chooses gradient with respect to given symbol.
  wrt :: Symbolic a => Gradients a -> Symbol a -> Symbol a
  wrt gradients x = HM.findWithDefault (symbolicZero x) x $ unGradients gradients
  
  -- | Takes partial gradients wrt to all symbols in a list sequentially, returning last result.
  nthPartialGradient :: Symbolic a => Symbol a -> [Symbol a] -> Symbol a
  nthPartialGradient = foldl' $ \y x -> getGradientsOf y `wrt` x
  
  -- | Takes nth order gradient of one symbol wrt other symbol. If n is negative number, an error is returned.
  nthGradient :: Symbolic a => Int -> Symbol a -> Symbol a -> Symbol a
  nthGradient n y x
      | n < 0 = error "Cannot take negative order gradient"
      | otherwise = nthPartialGradient y (replicate n x)
  