content: -----------------------------------------------------------------------------
  -- |
  -- Copyright   : (c) Edward Kmett 2010-2021
  -- License     : BSD3
  -- Maintainer  : ekmett@gmail.com
  -- Stability   : experimental
  -- Portability : GHC only
  --
  -- Higher order derivatives via a \"dual number tower\".
  --
  -----------------------------------------------------------------------------
  
  module Numeric.AD.Rank1.Sparse.Double
    ( SparseDouble
    , auto
    -- * Sparse Gradients
    , grad
    , grad'
    , gradWith
    , gradWith'
    -- * Variadic Gradients
    -- $vgrad
    , Grad
    , vgrad
    -- * Higher-Order Gradients
    , grads
    -- * Variadic Higher-Order Gradients
    , Grads
    , vgrads
  
    -- * Sparse Jacobians (synonyms)
    , jacobian
    , jacobian'
    , jacobianWith
    , jacobianWith'
    , jacobians
  
    -- * Sparse Hessians
    , hessian
    , hessian'
  
    , hessianF
    , hessianF'
  
    ) where
  
  import Control.Comonad
  import Control.Comonad.Cofree
  import Numeric.AD.Jet
  import Numeric.AD.Internal.Sparse.Double
  import Numeric.AD.Internal.Combinators
  import Numeric.AD.Mode
  
  second :: (a -> b) -> (c, a) -> (c, b)
  second g (a,b) = (a, g b)
  {-# INLINE second #-}
  
  grad
    :: Traversable f
    => (f SparseDouble -> SparseDouble)
    -> f Double -> f Double
  grad f as = d as $ apply f as
  {-# INLINE grad #-}
  
  grad'
    :: Traversable f
    => (f SparseDouble -> SparseDouble)
    -> f Double -> (Double, f Double)
  grad' f as = d' as $ apply f as
  {-# INLINE grad' #-}
  
  gradWith
    :: Traversable f
    => (Double -> Double -> b)
    -> (f SparseDouble -> SparseDouble)
    -> f Double
    -> f b
  gradWith g f as = zipWithT g as $ grad f as
  {-# INLINE gradWith #-}
  
  gradWith'
    :: Traversable f
    => (Double -> Double -> b)
    -> (f SparseDouble -> SparseDouble)
    -> f Double
    -> (Double, f b)
  gradWith' g f as = second (zipWithT g as) $ grad' f as
  {-# INLINE gradWith' #-}
  
  jacobian
    :: (Traversable f, Functor g)
    => (f SparseDouble -> g SparseDouble)
    -> f Double
    -> g (f Double)
  jacobian f as = d as <$> apply f as
  {-# INLINE jacobian #-}
  
  jacobian'
    :: (Traversable f, Functor g)
    => (f SparseDouble -> g SparseDouble)
    -> f Double
    -> g (Double, f Double)
  jacobian' f as = d' as <$> apply f as
  {-# INLINE jacobian' #-}
  
  jacobianWith
    :: (Traversable f, Functor g)
    => (Double -> Double -> b)
    -> (f SparseDouble -> g SparseDouble)
    -> f Double
    -> g (f b)
  jacobianWith g f as = zipWithT g as <$> jacobian f as
  {-# INLINE jacobianWith #-}
  
  jacobianWith'
    :: (Traversable f, Functor g)
    => (Double -> Double -> b)
    -> (f SparseDouble -> g SparseDouble)
    -> f Double
    -> g (Double, f b)
  jacobianWith' g f as = second (zipWithT g as) <$> jacobian' f as
  {-# INLINE jacobianWith' #-}
  
  grads
    :: Traversable f
    => (f SparseDouble -> SparseDouble)
    -> f Double
    -> Cofree f Double
  grads f as = ds as $ apply f as
  {-# INLINE grads #-}
  
  jacobians
    :: (Traversable f, Functor g)
    => (f SparseDouble -> g SparseDouble)
    -> f Double
    -> g (Cofree f Double)
  jacobians f as = ds as <$> apply f as
  {-# INLINE jacobians #-}
  
  d2 :: Functor f
    => Cofree f a
    -> f (f a)
  d2 = headJet . tailJet . tailJet . jet
  {-# INLINE d2 #-}
  
  d2'
    :: Functor f
    => Cofree f a
    -> (a, f (a, f a))
  d2' (a :< as) = (a, fmap (\(da :< das) -> (da, extract <$> das)) as)
  {-# INLINE d2' #-}
  
  hessian
    :: Traversable f
    => (f SparseDouble -> SparseDouble)
    -> f Double
    -> f (f Double)
  hessian f as = d2 $ grads f as
  {-# INLINE hessian #-}
  
  hessian'
    :: Traversable f
    => (f SparseDouble -> SparseDouble)
    -> f Double
    -> (Double, f (Double, f Double))
  hessian' f as = d2' $ grads f as
  {-# INLINE hessian' #-}
  
  hessianF
    :: (Traversable f, Functor g)
    => (f SparseDouble -> g SparseDouble)
    -> f Double
    -> g (f (f Double))
  hessianF f as = d2 <$> jacobians f as
  {-# INLINE hessianF #-}
  
  hessianF'
    :: (Traversable f, Functor g)
    => (f SparseDouble -> g SparseDouble)
    -> f Double
    -> g (Double, f (Double, f Double))
  hessianF' f as = d2' <$> jacobians f as
  {-# INLINE hessianF' #-}
  
  -- $vgrad
  --
  -- Variadic combinators for variadic mixed-mode automatic differentiation.
  --
  -- Unfortunately, variadicity comes at the expense of being able to use
  -- quantification to avoid sensitivity confusion, so be careful when
  -- counting the number of 'auto' calls you use when taking the gradient
  -- of a function that takes gradients!
  