{-# LANGUAGE RebindableSyntax #-}
{-# LANGUAGE TypeFamilies #-}

-- | Star-elimination evidence as a channel type.
--
-- The insight is that star-elimination structure is /data carried on a wire/,
-- not wiring itself.  @circuits@ keeps 'Circuit.Trace.yank' free of evidence;
-- 'circuits-diff' instead uses a channel whose values carry their own
-- elimination dictionary.
--
-- A @StarChannel s@ is both the feedback state (the @s@ field) and the
-- dictionary needed to solve affine feedback equations over that state.  The
-- dictionary part of the value is carried through the trace unchanged; the
-- data part is updated by the knot body.
module Circuit.Diff.Evidence
  ( -- * Evidence-carrying channel
    StarChannel (..),
    Scalar,

    -- * Concrete channels
    fieldStarChannel,
    listStarChannel,

    -- * Wrappers for knot bodies
    withStarChannel,
    withStarChannelDiff,
  )
where

import Circuit.Diff (Diff (..))
import Circuit.Mat.Dense (Matrix (..), fromLists, matVec, toLists)
import Circuit.Mat.Dense qualified as MD
import Circuit.Pullback (Pullback (..))
import NumHask.Algebra.Additive qualified as NHA
import NumHask.Algebra.Multiplicative qualified as NHM
import NumHask.Algebra.Ring qualified as NHR
import NumHask.Free.Carriers (FieldStar (..))
import NumHask.Prelude hiding (Scalar)

-- | Scalar carrier associated with a star-channel state type.
--
-- For scalar channels the scalar is the state itself; for list channels it
-- is the list element type.
type family Scalar s

type instance Scalar FieldStar = FieldStar

type instance Scalar [a] = a

-- | A feedback channel that carries its own star-elimination dictionary.
--
-- The matrix carrier is fixed by the 'Scalar' type family so that
-- @circuits-diff@ (which already depends on @circuits-mat@) can use dense
-- matrices directly, while @circuits@ remains ignorant of matrices.
--
-- A value bundles the current feedback state ('starData') with the operations
-- needed to eliminate a self-coupled affine knot over that state.  The
-- dictionary fields are expected to be preserved by a knot body;
-- 'withStarChannel' and 'withStarChannelDiff' are the canonical ways to build
-- such bodies.
data StarChannel s = StarChannel
  { -- | Dimension of the channel cotangent space.
    forall s. StarChannel s -> Int
starDim :: Int,
    -- | Current feedback state.
    forall s. StarChannel s -> s
starData :: s,
    -- | Zero channel cotangent of the given dimension.
    forall s. StarChannel s -> Int -> s
starZero :: Int -> s,
    -- | Basis vector @e_i@: @starBasis dim i@.
    forall s. StarChannel s -> Int -> Int -> s
starBasis :: Int -> Int -> s,
    -- | Componentwise sum of two channel cotangents.
    forall s. StarChannel s -> s -> s -> s
starAdd :: s -> s -> s,
    -- | Additive inverse of a channel cotangent.
    forall s. StarChannel s -> s -> s
starNegate :: s -> s,
    -- | Build the self-coupling matrix of a linear map @s -> s@ by
    -- probing each basis vector of the given dimension.
    forall s. StarChannel s -> Int -> (s -> s) -> Matrix (Scalar s)
starSelfMatrix :: Int -> (s -> s) -> Matrix (Scalar s),
    -- | Apply a matrix to a channel cotangent.
    forall s. StarChannel s -> Matrix (Scalar s) -> s -> s
starApplyMatrix :: Matrix (Scalar s) -> s -> s,
    -- | Kleene star of a square self-coupling matrix.
    forall s. StarChannel s -> Matrix (Scalar s) -> Matrix (Scalar s)
starMatrix :: Matrix (Scalar s) -> Matrix (Scalar s)
  }

-- | Evidence for a one-dimensional 'FieldStar' channel.
--
-- The matrix carrier is @'Matrix' 'FieldStar'@ — a 1×1 matrix whose single
-- element is the scalar channel cotangent.
fieldStarChannel :: StarChannel FieldStar
fieldStarChannel :: StarChannel FieldStar
fieldStarChannel =
  StarChannel
    { starDim :: Int
starDim = Int
1,
      starData :: FieldStar
starData = FieldStar
forall a. Additive a => a
NHA.zero,
      starZero :: Int -> FieldStar
starZero = FieldStar -> Int -> FieldStar
forall a b. a -> b -> a
const FieldStar
forall a. Additive a => a
NHA.zero,
      starBasis :: Int -> Int -> FieldStar
starBasis = \Int
_ Int
_ -> FieldStar
forall a. Multiplicative a => a
NHM.one,
      starAdd :: FieldStar -> FieldStar -> FieldStar
starAdd = FieldStar -> FieldStar -> FieldStar
forall a. Additive a => a -> a -> a
(NHA.+),
      starNegate :: FieldStar -> FieldStar
starNegate = FieldStar -> FieldStar
forall a. Subtractive a => a -> a
NHA.negate,
      starSelfMatrix :: Int -> (FieldStar -> FieldStar) -> Matrix (Scalar FieldStar)
starSelfMatrix = \Int
_ FieldStar -> FieldStar
f -> [[FieldStar]] -> Matrix FieldStar
forall a. [[a]] -> Matrix a
fromLists [[FieldStar -> FieldStar
f FieldStar
forall a. Multiplicative a => a
NHM.one]],
      starApplyMatrix :: Matrix (Scalar FieldStar) -> FieldStar -> FieldStar
starApplyMatrix = \Matrix (Scalar FieldStar)
m FieldStar
v -> case Matrix FieldStar -> [[FieldStar]]
forall a. Matrix a -> [[a]]
toLists Matrix FieldStar
Matrix (Scalar FieldStar)
m of
        [[FieldStar
s]] -> FieldStar
s FieldStar -> FieldStar -> FieldStar
forall a. Multiplicative a => a -> a -> a
NHM.* FieldStar
v
        [[FieldStar]]
_ -> [Char] -> FieldStar
forall a. HasCallStack => [Char] -> a
error [Char]
"Circuit.Diff.Evidence.applyMatrixE: scalar channel expected a 1x1 matrix",
      starMatrix :: Matrix (Scalar FieldStar) -> Matrix (Scalar FieldStar)
starMatrix = Matrix FieldStar -> Matrix FieldStar
Matrix (Scalar FieldStar) -> Matrix (Scalar FieldStar)
forall a. StarSemiring a => Matrix a -> Matrix a
MD.starMatrix
    }

-- | Evidence for an n-dimensional list channel.
--
-- The matrix carrier is @'Matrix' a@; the channel cotangent is @[a]@.
listStarChannel ::
  ( NHR.StarSemiring a,
    NHA.Subtractive a
  ) =>
  Int ->
  StarChannel [a]
listStarChannel :: forall a. (StarSemiring a, Subtractive a) => Int -> StarChannel [a]
listStarChannel Int
dim =
  let basisVec :: a -> a -> [a]
basisVec a
n a
i = [if a
k a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
i then a
forall a. Multiplicative a => a
NHM.one else a
forall a. Additive a => a
NHA.zero | a
k <- [a
0 .. a
n a -> a -> a
forall a. Subtractive a => a -> a -> a
- a
1]]
      zeroVec :: Int -> [a]
zeroVec Int
n = Int -> a -> [a]
forall a. Int -> a -> [a]
replicate Int
n a
forall a. Additive a => a
NHA.zero
   in StarChannel
        { starDim :: Int
starDim = Int
dim,
          starData :: [a]
starData = Int -> [a]
forall {a}. Additive a => Int -> [a]
zeroVec Int
dim,
          starZero :: Int -> [a]
starZero = Int -> [a]
forall {a}. Additive a => Int -> [a]
zeroVec,
          starBasis :: Int -> Int -> [a]
starBasis = Int -> Int -> [a]
forall {a} {a}.
(FromInteger a, Subtractive a, Enum a, Eq a, Multiplicative a,
 Additive a) =>
a -> a -> [a]
basisVec,
          starAdd :: [a] -> [a] -> [a]
starAdd = (a -> a -> a) -> [a] -> [a] -> [a]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith a -> a -> a
forall a. Additive a => a -> a -> a
(NHA.+),
          starNegate :: [a] -> [a]
starNegate = (a -> a) -> [a] -> [a]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> a
forall a. Subtractive a => a -> a
NHA.negate,
          starSelfMatrix :: Int -> ([a] -> [a]) -> Matrix (Scalar [a])
starSelfMatrix = \Int
n [a] -> [a]
f ->
            let cols :: [[a]]
cols = [[a] -> [a]
f (Int -> Int -> [a]
forall {a} {a}.
(FromInteger a, Subtractive a, Enum a, Eq a, Multiplicative a,
 Additive a) =>
a -> a -> [a]
basisVec Int
n Int
i) | Int
i <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Subtractive a => a -> a -> a
- Int
1]]
             in [[a]] -> Matrix a
forall a. [[a]] -> Matrix a
fromLists [[[a]
col [a] -> Int -> a
forall a. HasCallStack => [a] -> Int -> a
!! Int
k | [a]
col <- [[a]]
cols] | Int
k <- [Int
0 .. Int
n Int -> Int -> Int
forall a. Subtractive a => a -> a -> a
- Int
1]],
          starApplyMatrix :: Matrix (Scalar [a]) -> [a] -> [a]
starApplyMatrix = Matrix a -> [a] -> [a]
Matrix (Scalar [a]) -> [a] -> [a]
forall a. (Additive a, Multiplicative a) => Matrix a -> [a] -> [a]
matVec,
          starMatrix :: Matrix (Scalar [a]) -> Matrix (Scalar [a])
starMatrix = Matrix a -> Matrix a
Matrix (Scalar [a]) -> Matrix (Scalar [a])
forall a. StarSemiring a => Matrix a -> Matrix a
MD.starMatrix
        }

-- | Wrap a pullback computation on the underlying channel @s@ into a
-- computation on 'StarChannel' @s@.
--
-- The supplied @dict@ provides the elimination dictionary; the body only
-- needs to update 'starData'.  The output value reuses @dict@'s dictionary
-- fields, so 'Circuit.Diff.Star.solveStarBody' can read them directly from
-- the exposed channel type.
withStarChannel ::
  StarChannel s ->
  Pullback (s, b) (s, c) ->
  Pullback (StarChannel s, b) (StarChannel s, c)
withStarChannel :: forall s b c.
StarChannel s
-> Pullback (s, b) (s, c)
-> Pullback (StarChannel s, b) (StarChannel s, c)
withStarChannel StarChannel s
dict (Pullback (s, b) -> (s, c)
body) =
  ((StarChannel s, b) -> (StarChannel s, c))
-> Pullback (StarChannel s, b) (StarChannel s, c)
forall b a. (b -> a) -> Pullback b a
Pullback
    ( \(StarChannel s
sc, b
b) ->
        let (s
s', c
c) = (s, b) -> (s, c)
body (StarChannel s -> s
forall s. StarChannel s -> s
starData StarChannel s
sc, b
b)
         in (StarChannel s
dict {starData = s'}, c
c)
    )

-- | Wrap a 'Diff' computation on the underlying channel @s@ into a
-- computation on 'StarChannel' @s@.
--
-- This is the forward-differentiable analogue of 'withStarChannel': the
-- forward pass threads the 'StarChannel' state, and the backward pass
-- threads its cotangent through the same dictionary.
withStarChannelDiff ::
  StarChannel s ->
  Diff p (s, b) (s, c) ->
  Diff p (StarChannel s, b) (StarChannel s, c)
withStarChannelDiff :: forall {k} s (p :: k) b c.
StarChannel s
-> Diff p (s, b) (s, c)
-> Diff p (StarChannel s, b) (StarChannel s, c)
withStarChannelDiff StarChannel s
dict (Diff (s, b) -> ((s, c), (s, c) -> (s, b))
body) =
  ((StarChannel s, b)
 -> ((StarChannel s, c), (StarChannel s, c) -> (StarChannel s, b)))
-> Diff p (StarChannel s, b) (StarChannel s, c)
forall k (p :: k) a b. (a -> (b, b -> a)) -> Diff p a b
Diff
    ( \(StarChannel s
sc, b
b) ->
        let ((s
s', c
c), (s, c) -> (s, b)
back) = (s, b) -> ((s, c), (s, c) -> (s, b))
body (StarChannel s -> s
forall s. StarChannel s -> s
starData StarChannel s
sc, b
b)
            sc' :: StarChannel s
sc' = StarChannel s
dict {starData = s'}
            backward :: (StarChannel s, c) -> (StarChannel s, b)
backward (StarChannel s
dsc, c
dc) =
              let (s
ds, b
db) = (s, c) -> (s, b)
back (StarChannel s -> s
forall s. StarChannel s -> s
starData StarChannel s
dsc, c
dc)
               in (StarChannel s
dict {starData = ds}, b
db)
         in ((StarChannel s
sc', c
c), (StarChannel s, c) -> (StarChannel s, b)
backward)
    )