content: {-# LANGUAGE FlexibleContexts #-}
  {-# OPTIONS_HADDOCK not-home #-}
  
  -- |
  -- Module      : Prelude.Backprop.Explicit
  -- Copyright   : (c) Justin Le 2023
  -- License     : BSD3
  --
  -- Maintainer  : justin@jle.im
  -- Stability   : experimental
  -- Portability : non-portable
  --
  -- Provides "explicit" versions of all of the functions in
  -- "Prelude.Backprop".  Instead of relying on a 'Backprop' instance, allows
  -- you to manually provide 'zero', 'add', and 'one' on a per-value basis.
  --
  -- WARNING: API of this module can be considered only "semi-stable"; while
  -- the API of "Prelude.Backprop" and Prelude.Backprop.Num" are kept
  -- consistent, some argument order changes might happen in this module to
  -- reflect changes in underlying implementation.
  --
  -- @since 0.2.0.0
  module Prelude.Backprop.Explicit (
    -- * Foldable and Traversable
    sum,
    product,
    length,
    minimum,
    maximum,
    traverse,
    toList,
    mapAccumL,
    mapAccumR,
    foldr,
    foldl',
  
    -- * Functor and Applicative
    fmap,
    fmapConst,
    pure,
    liftA2,
    liftA3,
  
    -- * Numeric
    fromIntegral,
    realToFrac,
    round,
    fromIntegral',
  
    -- * Misc
    coerce,
  ) where
  
  import qualified Control.Applicative as P
  import Data.Bifunctor
  import qualified Data.Coerce as C
  import qualified Data.Foldable as P
  import qualified Data.Traversable as P
  import Numeric.Backprop.Explicit
  import Prelude (
    Applicative,
    Eq (..),
    Foldable,
    Fractional (..),
    Functor,
    Num (..),
    Ord (..),
    Traversable,
    ($),
    (.),
   )
  import qualified Prelude as P
  
  -- | 'Prelude.Backprop.sum', but taking explicit 'add' and 'zero'.
  sum ::
    (Foldable t, Functor t, Num a, Reifies s W) =>
    AddFunc (t a) ->
    BVar s (t a) ->
    BVar s a
  sum af = liftOp1 af . op1 $ \xs ->
    ( P.sum xs
    , (P.<$ xs)
    )
  {-# INLINE sum #-}
  
  -- | 'Prelude.Backprop.pure', but taking explicit 'add' and 'zero'.
  pure ::
    (Foldable t, Applicative t, Reifies s W) =>
    AddFunc a ->
    ZeroFunc a ->
    BVar s a ->
    BVar s (t a)
  pure af zfa = liftOp1 af . op1 $ \x ->
    ( P.pure x
    , \d -> case P.toList d of
        [] -> runZF zfa x
        e : es -> P.foldl' (runAF af) e es
    )
  {-# INLINE pure #-}
  
  -- | 'Prelude.Backprop.product', but taking explicit 'add' and 'zero'.
  product ::
    (Foldable t, Functor t, Fractional a, Reifies s W) =>
    AddFunc (t a) ->
    BVar s (t a) ->
    BVar s a
  product af = liftOp1 af . op1 $ \xs ->
    let p = P.product xs
     in ( p
        , \d -> (\x -> p * d / x) P.<$> xs
        )
  {-# INLINE product #-}
  
  -- | 'Prelude.Backprop.length', but taking explicit 'add' and 'zero'.
  length ::
    (Foldable t, Num b, Reifies s W) =>
    AddFunc (t a) ->
    ZeroFunc (t a) ->
    BVar s (t a) ->
    BVar s b
  length af zfa = liftOp1 af . op1 $ \xs ->
    ( P.fromIntegral (P.length xs)
    , P.const (runZF zfa xs)
    )
  {-# INLINE length #-}
  
  -- | 'Prelude.Backprop.minimum', but taking explicit 'add' and 'zero'.
  minimum ::
    (Foldable t, Functor t, Ord a, Reifies s W) =>
    AddFunc (t a) ->
    ZeroFunc a ->
    BVar s (t a) ->
    BVar s a
  minimum af zf = liftOp1 af . op1 $ \xs ->
    let m = P.minimum xs
     in ( m
        , \d -> (\x -> if x == m then d else runZF zf x) P.<$> xs
        )
  {-# INLINE minimum #-}
  
  -- | 'Prelude.Backprop.maximum', but taking explicit 'add' and 'zero'.
  maximum ::
    (Foldable t, Functor t, Ord a, Reifies s W) =>
    AddFunc (t a) ->
    ZeroFunc a ->
    BVar s (t a) ->
    BVar s a
  maximum af zf = liftOp1 af . op1 $ \xs ->
    let m = P.maximum xs
     in ( m
        , \d -> (\x -> if x == m then d else runZF zf x) P.<$> xs
        )
  {-# INLINE maximum #-}
  
  -- | 'Prelude.Backprop.foldr', but taking explicit 'add' and 'zero'.
  --
  -- @since 0.2.3.0
  foldr ::
    (Traversable t, Reifies s W) =>
    AddFunc a ->
    ZeroFunc a ->
    (BVar s a -> BVar s b -> BVar s b) ->
    BVar s b ->
    BVar s (t a) ->
    BVar s b
  foldr af z f x = P.foldr f x . toList af z
  {-# INLINE foldr #-}
  
  -- | 'Prelude.Backprop.foldl'', but taking explicit 'add' and 'zero'.
  --
  -- @since 0.2.3.0
  foldl' ::
    (Traversable t, Reifies s W) =>
    AddFunc a ->
    ZeroFunc a ->
    (BVar s b -> BVar s a -> BVar s b) ->
    BVar s b ->
    BVar s (t a) ->
    BVar s b
  foldl' af z f x = P.foldl' f x . toList af z
  {-# INLINE foldl' #-}
  
  -- | 'Prelude.Backprop.fmap', but taking explicit 'add' and 'zero'.
  fmap ::
    (Traversable f, Reifies s W) =>
    AddFunc a ->
    AddFunc b ->
    ZeroFunc a ->
    ZeroFunc b ->
    (BVar s a -> BVar s b) ->
    BVar s (f a) ->
    BVar s (f b)
  fmap afa afb zfa zfb f = collectVar afb zfb . P.fmap f . sequenceVar afa zfa
  {-# INLINE fmap #-}
  
  -- | 'Prelude.Backprop.fmapConst', but taking explicit 'add' and 'zero'.
  --
  -- @since 0.2.4.0
  fmapConst ::
    (Functor f, Foldable f, Reifies s W) =>
    AddFunc (f a) ->
    AddFunc b ->
    ZeroFunc (f a) ->
    ZeroFunc b ->
    BVar s b ->
    BVar s (f a) ->
    BVar s (f b)
  fmapConst afa afb zfa zfb = liftOp2 afb afa . op2 $ \x xs ->
    ( x P.<$ xs
    , \d ->
        ( case P.toList d of
            [] -> runZF zfb x
            e : es -> P.foldl' (runAF afb) e es
        , runZF zfa xs
        )
    )
  {-# INLINE fmapConst #-}
  
  -- | 'Prelude.Backprop.traverse', but taking explicit 'add' and 'zero'.
  traverse ::
    (Traversable t, Applicative f, Foldable f, Reifies s W) =>
    AddFunc a ->
    AddFunc b ->
    AddFunc (t b) ->
    ZeroFunc a ->
    ZeroFunc b ->
    (BVar s a -> f (BVar s b)) ->
    BVar s (t a) ->
    BVar s (f (t b))
  traverse afa afb aftb zfa zfb f =
    collectVar aftb zftb
      . P.fmap (collectVar afb zfb)
      . P.traverse f
      . sequenceVar afa zfa
    where
      zftb = ZF $ P.fmap (runZF zfb)
      {-# INLINE zftb #-}
  {-# INLINE traverse #-}
  
  -- | 'Prelude.Backprop.liftA2', but taking explicit 'add' and 'zero'.
  liftA2 ::
    ( Traversable f
    , Applicative f
    , Reifies s W
    ) =>
    AddFunc a ->
    AddFunc b ->
    AddFunc c ->
    ZeroFunc a ->
    ZeroFunc b ->
    ZeroFunc c ->
    (BVar s a -> BVar s b -> BVar s c) ->
    BVar s (f a) ->
    BVar s (f b) ->
    BVar s (f c)
  liftA2 afa afb afc zfa zfb zfc f x y =
    collectVar afc zfc $
      f
        P.<$> sequenceVar afa zfa x
        P.<*> sequenceVar afb zfb y
  {-# INLINE liftA2 #-}
  
  -- | 'Prelude.Backprop.liftA3', but taking explicit 'add' and 'zero'.
  liftA3 ::
    ( Traversable f
    , Applicative f
    , Reifies s W
    ) =>
    AddFunc a ->
    AddFunc b ->
    AddFunc c ->
    AddFunc d ->
    ZeroFunc a ->
    ZeroFunc b ->
    ZeroFunc c ->
    ZeroFunc d ->
    (BVar s a -> BVar s b -> BVar s c -> BVar s d) ->
    BVar s (f a) ->
    BVar s (f b) ->
    BVar s (f c) ->
    BVar s (f d)
  liftA3 afa afb afc afd zfa zfb zfc zfd f x y z =
    collectVar afd zfd $
      f
        P.<$> sequenceVar afa zfa x
        P.<*> sequenceVar afb zfb y
        P.<*> sequenceVar afc zfc z
  {-# INLINE liftA3 #-}
  
  -- | Coerce items inside a 'BVar'.
  coerce :: C.Coercible a b => BVar s a -> BVar s b
  coerce = coerceVar
  {-# INLINE coerce #-}
  
  -- | 'Prelude.Backprop.fromIntegral', but taking explicit 'add' and 'zero'.
  --
  -- @since 0.2.1.0
  fromIntegral ::
    (P.Integral a, P.Integral b, Reifies s W) =>
    AddFunc a ->
    BVar s a ->
    BVar s b
  fromIntegral af = isoVar af P.fromIntegral P.fromIntegral
  {-# INLINE fromIntegral #-}
  
  -- | 'Prelude.Backprop.realToFrac', but taking explicit 'add' and 'zero'.
  --
  -- @since 0.2.1.0
  realToFrac ::
    (Fractional a, P.Real a, Fractional b, P.Real b, Reifies s W) =>
    AddFunc a ->
    BVar s a ->
    BVar s b
  realToFrac af = isoVar af P.realToFrac P.realToFrac
  {-# INLINE realToFrac #-}
  
  -- | 'Prelude.Backprop.round', but taking explicit 'add' and 'zero'.
  --
  -- @since 0.2.3.0
  round ::
    (P.RealFrac a, P.Integral b, Reifies s W) =>
    AddFunc a ->
    BVar s a ->
    BVar s b
  round af = isoVar af P.round P.fromIntegral
  {-# INLINE round #-}
  
  -- | 'Prelude.Backprop.fromIntegral'', but taking explicit 'add' and
  -- 'zero'.
  --
  -- @since 0.2.3.0
  fromIntegral' ::
    (P.Integral a, P.RealFrac b, Reifies s W) =>
    AddFunc a ->
    BVar s a ->
    BVar s b
  fromIntegral' af = isoVar af P.fromIntegral P.round
  {-# INLINE fromIntegral' #-}
  
  -- | 'Prelude.Backprop.length', but taking explicit 'add' and 'zero'.
  --
  -- @since 0.2.2.0
  toList ::
    (Traversable t, Reifies s W) =>
    AddFunc a ->
    ZeroFunc a ->
    BVar s (t a) ->
    [BVar s a]
  toList af z = toListOfVar af (ZF (P.fmap (runZF z))) P.traverse
  {-# INLINE toList #-}
  
  -- | 'Prelude.Backprop.mapAccumL', but taking explicit 'add' and 'zero'.
  --
  -- @since 0.2.2.0
  mapAccumL ::
    (Traversable t, Reifies s W) =>
    AddFunc b ->
    AddFunc c ->
    ZeroFunc b ->
    ZeroFunc c ->
    (BVar s a -> BVar s b -> (BVar s a, BVar s c)) ->
    BVar s a ->
    BVar s (t b) ->
    (BVar s a, BVar s (t c))
  mapAccumL afb afc zfb zfc f s =
    second (collectVar afc zfc)
      . P.mapAccumL f s
      . sequenceVar afb zfb
  {-# INLINE mapAccumL #-}
  
  -- | 'Prelude.Backprop.mapAccumR', but taking explicit 'add' and 'zero'.
  --
  -- @since 0.2.2.0
  mapAccumR ::
    (Traversable t, Reifies s W) =>
    AddFunc b ->
    AddFunc c ->
    ZeroFunc b ->
    ZeroFunc c ->
    (BVar s a -> BVar s b -> (BVar s a, BVar s c)) ->
    BVar s a ->
    BVar s (t b) ->
    (BVar s a, BVar s (t c))
  mapAccumR afb afc zfb zfc f s =
    second (collectVar afc zfc)
      . P.mapAccumR f s
      . sequenceVar afb zfb
  {-# INLINE mapAccumR #-}
  