content: {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE EmptyCase #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE FunctionalDependencies #-}
  {-# LANGUAGE GADTs #-}
  {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE RankNTypes #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TypeApplications #-}
  {-# LANGUAGE TypeOperators #-}
  {-# LANGUAGE UndecidableInstances #-}
  {-# OPTIONS_HADDOCK not-home #-}
  
  -- |
  -- Module      : Numeric.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
  -- "Numeric.Backprop".  Instead of relying on a 'Backprop' instance, allows
  -- you to manually provide 'zero', 'add', and 'one' on a per-value basis.
  --
  -- It is recommended you use "Numeric.Backprop" or "Numeric.Backprop.Num"
  -- instead, unless your type has no 'Num' instance, or you else you want to
  -- avoid defining orphan 'Backprop' instances for external types.  Can also
  -- be useful if mixing and matching styles.
  --
  -- See "Numeric.Backprop" for fuller documentation on using these
  -- functions.
  --
  -- WARNING: API of this module can be considered only "semi-stable"; while
  -- the API of "Numeric.Backprop" and "Numeric.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 Numeric.Backprop.Explicit (
    -- * Types
    BVar,
    W,
    Backprop (..),
    ABP (..),
    NumBP (..),
  
    -- * Explicit 'zero', 'add', and 'one'
    ZeroFunc (..),
    zfNum,
    zfNums,
    zeroFunc,
    zeroFuncs,
    zfFunctor,
    AddFunc (..),
    afNum,
    afNums,
    addFunc,
    addFuncs,
    OneFunc (..),
    ofNum,
    ofNums,
    oneFunc,
    oneFuncs,
    ofFunctor,
  
    -- * Running
    backprop,
    evalBP,
    gradBP,
    backpropWith,
  
    -- ** Multiple inputs
    evalBP0,
    backprop2,
    evalBP2,
    gradBP2,
    backpropWith2,
    backpropN,
    evalBPN,
    gradBPN,
    backpropWithN,
    RPureConstrained,
  
    -- * Manipulating 'BVar'
    constVar,
    auto,
    coerceVar,
    viewVar,
    setVar,
    overVar,
    sequenceVar,
    collectVar,
    previewVar,
    toListOfVar,
  
    -- ** With Isomorphisms
    isoVar,
    isoVar2,
    isoVar3,
    isoVarN,
  
    -- ** With 'Op's
    liftOp,
    liftOp1,
    liftOp2,
    liftOp3,
  
    -- ** Generics
    splitBV,
    joinBV,
    BVGroup,
  
    -- * 'Op'
    Op (..),
  
    -- ** Creation
    op0,
    opConst,
    idOp,
    bpOp,
  
    -- *** Giving gradients directly
    op1,
    op2,
    op3,
  
    -- *** From Isomorphisms
    opCoerce,
    opTup,
    opIso,
    opIsoN,
    opLens,
  
    -- *** No gradients
    noGrad1,
    noGrad,
  
    -- * Utility
    Reifies,
  ) where
  
  import Data.Bifunctor
  import Data.Functor.Identity
  import Data.Reflection
  import Data.Type.Util
  import Data.Vinyl.Core
  import Data.Vinyl.TypeLevel
  import GHC.Generics as G
  import Lens.Micro
  import Numeric.Backprop.Class
  import Numeric.Backprop.Internal
  import Numeric.Backprop.Op
  import Unsafe.Coerce
  
  -- | 'ZeroFunc's for every item in a type level list based on their
  -- 'Num' instances
  --
  -- @since 0.2.0.0
  zfNums :: RPureConstrained Num as => Rec ZeroFunc as
  zfNums = rpureConstrained @Num zfNum
  
  -- | 'zeroFunc' for instances of 'Functor'
  --
  -- @since 0.2.1.0
  zfFunctor :: (Backprop a, Functor f) => ZeroFunc (f a)
  zfFunctor = ZF zeroFunctor
  {-# INLINE zfFunctor #-}
  
  -- | 'ZeroFunc's for every item in a type level list based on their
  -- 'Num' instances
  --
  -- @since 0.2.0.0
  afNums :: RPureConstrained Num as => Rec AddFunc as
  afNums = rpureConstrained @Num afNum
  
  -- | 'ZeroFunc's for every item in a type level list based on their
  -- 'Num' instances
  --
  -- @since 0.2.0.0
  ofNums :: RPureConstrained Num as => Rec OneFunc as
  ofNums = rpureConstrained @Num ofNum
  
  -- | 'OneFunc' for instances of 'Functor'
  --
  -- @since 0.2.1.0
  ofFunctor :: (Backprop a, Functor f) => OneFunc (f a)
  ofFunctor = OF oneFunctor
  {-# INLINE ofFunctor #-}
  
  -- | Generate an 'ZeroFunc' for every type in a type-level list, if every
  -- type has an instance of 'Backprop'.
  --
  -- @since 0.2.0.0
  zeroFuncs :: RPureConstrained Backprop as => Rec ZeroFunc as
  zeroFuncs = rpureConstrained @Backprop zeroFunc
  
  -- | Generate an 'AddFunc' for every type in a type-level list, if every
  -- type has an instance of 'Backprop'.
  --
  -- @since 0.2.0.0
  addFuncs :: RPureConstrained Backprop as => Rec AddFunc as
  addFuncs = rpureConstrained @Backprop addFunc
  
  -- | Generate an 'OneFunc' for every type in a type-level list, if every
  -- type has an instance of 'Backprop'.
  --
  -- @since 0.2.0.0
  oneFuncs :: RPureConstrained Backprop as => Rec OneFunc as
  oneFuncs = rpureConstrained @Backprop oneFunc
  
  -- | Shorter alias for 'constVar', inspired by the /ad/ library.
  --
  -- @since 0.2.0.0
  auto :: a -> BVar s a
  auto = constVar
  {-# INLINE auto #-}
  
  -- | 'Numeric.Backprop.backpropN', but with explicit 'zero' and 'one'.
  backpropN ::
    forall as b.
    () =>
    Rec ZeroFunc as ->
    OneFunc b ->
    (forall s. Reifies s W => Rec (BVar s) as -> BVar s b) ->
    Rec Identity as ->
    (b, Rec Identity as)
  backpropN zfs ob f xs = case backpropWithN zfs f xs of
    (y, g) -> (y, g (runOF ob y))
  {-# INLINE backpropN #-}
  
  -- | 'Numeric.Backprop.backprop', but with explicit 'zero' and 'one'.
  backprop ::
    ZeroFunc a ->
    OneFunc b ->
    (forall s. Reifies s W => BVar s a -> BVar s b) ->
    a ->
    (b, a)
  backprop zfa ofb f =
    second (\case Identity x :& RNil -> x)
      . backpropN (zfa :& RNil) ofb (f . (\case x :& RNil -> x))
      . (:& RNil)
      . Identity
  {-# INLINE backprop #-}
  
  -- | 'Numeric.Backprop.backpropWith', but with explicit 'zero'.
  --
  -- Note that argument order changed in v0.2.4.
  backpropWith ::
    ZeroFunc a ->
    (forall s. Reifies s W => BVar s a -> BVar s b) ->
    a ->
    (b, b -> a)
  backpropWith zfa f =
    second ((\case Identity x :& RNil -> x) .)
      . backpropWithN (zfa :& RNil) (f . (\case x :& RNil -> x))
      . (:& RNil)
      . Identity
  {-# INLINE backpropWith #-}
  
  -- | 'evalBP' but with no arguments.  Useful when everything is just given
  -- through 'constVar'.
  evalBP0 :: (forall s. Reifies s W => BVar s a) -> a
  evalBP0 x = evalBPN (const x) RNil
  {-# INLINE evalBP0 #-}
  
  -- | Turn a function @'BVar' s a -> 'BVar' s b@ into the function @a -> b@
  -- that it represents.
  --
  -- Benchmarks show that this should have virtually no overhead over
  -- directly writing a @a -> b@. 'BVar' is, in this situation, a zero-cost
  -- abstraction, performance-wise.
  --
  -- See documentation of 'Numeric.Backprop.backprop' for more information.
  evalBP :: (forall s. Reifies s W => BVar s a -> BVar s b) -> a -> b
  evalBP f = evalBPN (f . (\case x :& RNil -> x)) . (:& RNil) . Identity
  {-# INLINE evalBP #-}
  
  -- | 'Numeric.Backprop.gradBP', but with explicit 'zero' and 'one'.
  gradBP ::
    ZeroFunc a ->
    OneFunc b ->
    (forall s. Reifies s W => BVar s a -> BVar s b) ->
    a ->
    a
  gradBP zfa ofb f = snd . backprop zfa ofb f
  {-# INLINE gradBP #-}
  
  -- | 'Numeric.Backprop.gradBP', Nbut with explicit 'zero' and 'one'.
  gradBPN ::
    Rec ZeroFunc as ->
    OneFunc b ->
    (forall s. Reifies s W => Rec (BVar s) as -> BVar s b) ->
    Rec Identity as ->
    Rec Identity as
  gradBPN zfas ofb f = snd . backpropN zfas ofb f
  {-# INLINE gradBPN #-}
  
  -- | 'Numeric.Backprop.backprop2', but with explicit 'zero' and 'one'.
  backprop2 ::
    ZeroFunc a ->
    ZeroFunc b ->
    OneFunc c ->
    (forall s. Reifies s W => BVar s a -> BVar s b -> BVar s c) ->
    a ->
    b ->
    (c, (a, b))
  backprop2 zfa zfb ofc f x y =
    second (\(Identity dx :& Identity dy :& RNil) -> (dx, dy)) $
      backpropN
        (zfa :& zfb :& RNil)
        ofc
        (\(x' :& y' :& RNil) -> f x' y')
        (Identity x :& Identity y :& RNil)
  {-# INLINE backprop2 #-}
  
  -- | 'Numeric.Backprop.backpropWith2', but with explicit 'zero'.
  --
  -- Note that argument order changed in v0.2.4.
  --
  -- @since 0.2.0.0
  backpropWith2 ::
    ZeroFunc a ->
    ZeroFunc b ->
    (forall s. Reifies s W => BVar s a -> BVar s b -> BVar s c) ->
    a ->
    b ->
    (c, c -> (a, b))
  backpropWith2 zfa zfb f x y =
    second ((\(Identity dx :& Identity dy :& RNil) -> (dx, dy)) .) $
      backpropWithN
        (zfa :& zfb :& RNil)
        (\(x' :& y' :& RNil) -> f x' y')
        (Identity x :& Identity y :& RNil)
  {-# INLINE backpropWith2 #-}
  
  -- | 'evalBP' for a two-argument function.  See
  -- 'Numeric.Backprop.backprop2' for notes.
  evalBP2 ::
    (forall s. Reifies s W => BVar s a -> BVar s b -> BVar s c) ->
    a ->
    b ->
    c
  evalBP2 f x y =
    evalBPN (\(x' :& y' :& RNil) -> f x' y') $
      Identity x
        :& Identity y
        :& RNil
  {-# INLINE evalBP2 #-}
  
  -- | 'Numeric.Backprop.gradBP2' with explicit 'zero' and 'one'.
  gradBP2 ::
    ZeroFunc a ->
    ZeroFunc b ->
    OneFunc c ->
    (forall s. Reifies s W => BVar s a -> BVar s b -> BVar s c) ->
    a ->
    b ->
    (a, b)
  gradBP2 zfa zfb ofc f x = snd . backprop2 zfa zfb ofc f x
  {-# INLINE gradBP2 #-}
  
  -- | 'Numeric.Backprop.bpOp' with explicit 'zero'.
  bpOp ::
    Rec ZeroFunc as ->
    (forall s. Reifies s W => Rec (BVar s) as -> BVar s b) ->
    Op as b
  bpOp zfs f = Op (backpropWithN zfs f)
  {-# INLINE bpOp #-}
  
  -- | 'Numeric.Backprop.overVar' with explicit 'add' and 'zero'.
  --
  -- @since 0.2.4.0
  overVar ::
    Reifies s W =>
    AddFunc a ->
    AddFunc b ->
    ZeroFunc a ->
    ZeroFunc b ->
    Lens' b a ->
    (BVar s a -> BVar s a) ->
    BVar s b ->
    BVar s b
  overVar afa afb zfa zfb l f x = setVar afa afb zfa l (f (viewVar afa zfb l x)) x
  {-# INLINE overVar #-}
  
  -- | 'Numeric.Backprop.isoVar' with explicit 'add' and 'zero'.
  isoVar ::
    Reifies s W =>
    AddFunc a ->
    (a -> b) ->
    (b -> a) ->
    BVar s a ->
    BVar s b
  isoVar af f g = liftOp1 af (opIso f g)
  {-# INLINE isoVar #-}
  
  -- | 'Numeric.Backprop.isoVar2' with explicit 'add' and 'zero'.
  isoVar2 ::
    Reifies s W =>
    AddFunc a ->
    AddFunc b ->
    (a -> b -> c) ->
    (c -> (a, b)) ->
    BVar s a ->
    BVar s b ->
    BVar s c
  isoVar2 afa afb f g = liftOp2 afa afb (opIso2 f g)
  {-# INLINE isoVar2 #-}
  
  -- | 'Numeric.Backprop.isoVar3' with explicit 'add' and 'zero'.
  isoVar3 ::
    Reifies s W =>
    AddFunc a ->
    AddFunc b ->
    AddFunc c ->
    (a -> b -> c -> d) ->
    (d -> (a, b, c)) ->
    BVar s a ->
    BVar s b ->
    BVar s c ->
    BVar s d
  isoVar3 afa afb afc f g = liftOp3 afa afb afc (opIso3 f g)
  {-# INLINE isoVar3 #-}
  
  -- | 'Numeric.Backprop.isoVarN' with explicit 'add' and 'zero'.
  isoVarN ::
    Reifies s W =>
    Rec AddFunc as ->
    (Rec Identity as -> b) ->
    (b -> Rec Identity as) ->
    Rec (BVar s) as ->
    BVar s b
  isoVarN afs f g = liftOp afs (opIsoN f g)
  {-# INLINE isoVarN #-}
  
  -- | Helper class for generically "splitting" and "joining" 'BVar's into
  -- constructors.  See 'Numeric.Backprop.splitBV' and
  -- 'Numeric.Backprop.joinBV'.
  --
  -- See "Numeric.Backprop#hkd" for a tutorial on how to use this.
  --
  -- Instances should be available for types made with one constructor whose
  -- fields are all instances of 'Backprop', with a 'Generic' instance.
  --
  -- @since 0.2.2.0
  class BVGroup s as i o | o -> i, i -> as where
    -- | Helper method for generically "splitting" 'BVar's out of
    -- constructors inside a 'BVar'.  See 'splitBV'.
    gsplitBV :: Rec AddFunc as -> Rec ZeroFunc as -> BVar s (i ()) -> o ()
  
    -- | Helper method for generically "joining" 'BVar's inside
    -- a constructor into a 'BVar'.  See 'joinBV'.
    gjoinBV :: Rec AddFunc as -> Rec ZeroFunc as -> o () -> BVar s (i ())
  
  instance BVGroup s '[] (K1 i a) (K1 i (BVar s a)) where
    gsplitBV _ _ = K1 . coerceVar
    {-# INLINE gsplitBV #-}
    gjoinBV _ _ = coerceVar . unK1
    {-# INLINE gjoinBV #-}
  
  instance
    BVGroup s as i o =>
    BVGroup s as (M1 p c i) (M1 p c o)
    where
    gsplitBV afs zfs = M1 . gsplitBV afs zfs . coerceVar @_ @(i ())
    {-# INLINE gsplitBV #-}
    gjoinBV afs zfs = coerceVar @(i ()) . gjoinBV afs zfs . unM1
    {-# INLINE gjoinBV #-}
  
  instance BVGroup s '[] V1 V1 where
    gsplitBV _ _ = unsafeCoerce
    {-# INLINE gsplitBV #-}
    gjoinBV _ _ = \case {}
    {-# INLINE gjoinBV #-}
  
  instance BVGroup s '[] U1 U1 where
    gsplitBV _ _ _ = U1
    {-# INLINE gsplitBV #-}
    gjoinBV _ _ _ = constVar U1
    {-# INLINE gjoinBV #-}
  
  instance
    ( Reifies s W
    , BVGroup s as i1 o1
    , BVGroup s bs i2 o2
    , cs ~ (as ++ bs)
    , RecApplicative as
    ) =>
    BVGroup s (i1 () ': i2 () ': cs) (i1 :*: i2) (o1 :*: o2)
    where
    gsplitBV (afa :& afb :& afs) (zfa :& zfb :& zfs) xy = x :*: y
      where
        (afas, afbs) = splitRec afs
        (zfas, zfbs) = splitRec zfs
        zfab = ZF $ \(xx :*: yy) -> runZF zfa xx :*: runZF zfb yy
        x = gsplitBV afas zfas . viewVar afa zfab p1 $ xy
        y = gsplitBV afbs zfbs . viewVar afb zfab p2 $ xy
    {-# INLINE gsplitBV #-}
    gjoinBV (afa :& afb :& afs) (_ :& _ :& zfs) (x :*: y) =
      isoVar2
        afa
        afb
        (:*:)
        unP
        (gjoinBV afas zfas x)
        (gjoinBV afbs zfbs y)
      where
        (afas, afbs) = splitRec afs
        (zfas, zfbs) = splitRec zfs
        unP (xx :*: yy) = (xx, yy)
    {-# INLINE gjoinBV #-}
  
  -- | This instance is possible but it is not clear when it would be useful
  instance
    ( Reifies s W
    , BVGroup s as i1 o1
    , BVGroup s bs i2 o2
    , cs ~ (as ++ bs)
    , RecApplicative as
    ) =>
    BVGroup s (i1 () ': i2 () ': cs) (i1 :+: i2) (o1 :+: o2)
    where
    gsplitBV (afa :& afb :& afs) (zfa :& zfb :& zfs) xy =
      case previewVar afa zf s1 xy of
        Just x -> L1 $ gsplitBV afas zfas x
        Nothing -> case previewVar afb zf s2 xy of
          Just y -> R1 $ gsplitBV afbs zfbs y
          Nothing -> error "Numeric.Backprop.gsplitBV: Internal error occurred"
      where
        zf = ZF $ \case
          L1 xx -> L1 $ runZF zfa xx
          R1 yy -> R1 $ runZF zfb yy
        (afas, afbs) = splitRec afs
        (zfas, zfbs) = splitRec zfs
    {-# INLINE gsplitBV #-}
    gjoinBV (afa :& afb :& afs) (zfa :& zfb :& zfs) = \case
      L1 x ->
        liftOp1
          afa
          (op1 (\xx -> (L1 xx, \case L1 d -> d; R1 _ -> runZF zfa xx)))
          (gjoinBV afas zfas x)
      R1 y ->
        liftOp1
          afb
          (op1 (\yy -> (R1 yy, \case L1 _ -> runZF zfb yy; R1 d -> d)))
          (gjoinBV afbs zfbs y)
      where
        (afas, afbs) = splitRec afs
        (zfas, zfbs) = splitRec zfs
    {-# INLINE gjoinBV #-}
  
  -- | 'Numeric.Backprop.splitBV' with explicit 'add' and 'zero'.
  --
  -- @since 0.2.2.0
  splitBV ::
    forall z f s as.
    ( Generic (z f)
    , Generic (z (BVar s))
    , BVGroup s as (Rep (z f)) (Rep (z (BVar s)))
    , Reifies s W
    ) =>
    AddFunc (Rep (z f) ()) ->
    Rec AddFunc as ->
    ZeroFunc (z f) ->
    Rec ZeroFunc as ->
    -- | 'BVar' of value
    BVar s (z f) ->
    -- | 'BVar's of fields
    z (BVar s)
  splitBV af afs zf zfs =
    G.to
      . gsplitBV afs zfs
      . viewVar af zf (lens (from @(z f) @()) (const G.to))
  {-# INLINE splitBV #-}
  
  -- | 'Numeric.Backprop.joinBV' with explicit 'add' and 'zero'.
  --
  -- @since 0.2.2.0
  joinBV ::
    forall z f s as.
    ( Generic (z f)
    , Generic (z (BVar s))
    , BVGroup s as (Rep (z f)) (Rep (z (BVar s)))
    , Reifies s W
    ) =>
    AddFunc (z f) ->
    Rec AddFunc as ->
    ZeroFunc (Rep (z f) ()) ->
    Rec ZeroFunc as ->
    -- | 'BVar's of fields
    z (BVar s) ->
    -- | 'BVar' of combined value
    BVar s (z f)
  joinBV af afs zf zfs =
    viewVar af zf (lens G.to (const from))
      . gjoinBV afs zfs
      . from @(z (BVar s)) @()
  {-# INLINE joinBV #-}
  