{-# LANGUAGE RebindableSyntax #-}

-- | The Schur-complement bridge — 'trace' solved by 'starMatrix'.
--
-- This is where the three floors actually touch.  A 'Diff knot body
-- has a backward pass that is /affine in cotangents/ (calculus
-- promises linearity), so its channel self-coupling is a linear map
-- recoverable by probing with basis cotangents.  Build that map as a
-- 'Matrix', take its Kleene star with 'starMatrix' — Gaussian
-- elimination, @(I − A)⁻¹@, the fourth face of the four-for-one — and
-- the trace's pullback is the Schur complement
--
-- > db = D·dc + B · star A · C·dc
--
-- recovered with one final probe at the backward fixpoint.
--
-- Cost: @dim@ probes of @backward@ plus one @starMatrix@ per forward
-- point, /shared across all cotangents/; each subsequent pullback is
-- two probes and a matrix–vector product.  Compare 'Circuit.Diff.Circuit.traceNFrom',
-- which pays @n@ probes per cotangent and is only as exact as @n@ is
-- large.  Here the backward pass is exact whenever the star exists.
--
-- The dependency is deliberate: this module imports "Circuit.Mat.Dense",
-- making circuits-ad ⇄ star-matrix a literal edge rather than a nominal
-- alias.
module Circuit.Diff.Star
  ( -- * Polymorphic bridge
    traceStarMatrix,

    -- * Body bridge
    solveStarBody,

    -- * Double adapters (via 'FieldStar')
    traceStarFromD,
    traceStarMatrixD,
  )
where

import Circuit.Bimonoid (Merge, MergeZero, Zero)
import Circuit.Bimonoid qualified as CB
import Circuit.Body (Body (..))
import Circuit.Diff.Circuit (Diff (..), traceStarFrom)
import Circuit.Diff.Evidence (StarChannel (..))
import Circuit.Mat.Dense (fromLists, matVec)
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

-- $setup
-- >>> import Circuit.Body (Body (..))
-- >>> import Circuit.Diff.Circuit
-- >>> import Circuit.Diff.Evidence (StarChannel (..), fieldStarChannel, withStarChannel)
-- >>> import Circuit.Pullback (Pullback (..))
-- >>> import NumHask.Free.Carriers (FieldStar (..))
-- >>> import Prelude hiding (id, (.))

-- | Trace over a /vector/ feedback channel, backward pass solved by
-- 'starMatrix'.
--
-- The channel is a list @[j]@ whose dimension is fixed by the seed.
-- Forward iterates from the seed (no closed form for a nonlinear
-- fixpoint).  Backward:
--
--   1. probe @backward@ with each basis cotangent @e_i@ to read off
--      column @i@ of the channel self-coupling A (in cotangent space);
--   2. @star A@ by Conway block recursion — 'starMatrix';
--   3. per cotangent @dc@: @C·dc = fst (backward (0, dc))@, then
--      @db = snd (backward (star A · C·dc, dc))@.
--
-- A and @star A@ are forced once (lazily, on first pullback) and
-- shared across cotangents.
--
-- __Proof obligation__: probing assumes every primitive's pullback is
-- a genuinely linear map — true for honestly-constructed 'Diff prims.
traceStarMatrix ::
  (NHR.StarSemiring j, MergeZero (->) c) =>
  -- | forward seed; its length is the channel dimension
  [j] ->
  -- | forward iteration count
  Int ->
  Diff p ([j], b) ([j], c) ->
  Diff p b c
traceStarMatrix :: forall {k} j c (p :: k) b.
(StarSemiring j, MergeZero (->) c) =>
[j] -> Int -> Diff p ([j], b) ([j], c) -> Diff p b c
traceStarMatrix [j]
x0 Int
n (Diff ([j], b) -> (([j], c), ([j], c) -> ([j], b))
body) = (b -> (c, c -> b)) -> Diff p b c
forall k (p :: k) a b. (a -> (b, b -> a)) -> Diff p a b
Diff ((b -> (c, c -> b)) -> Diff p b c)
-> (b -> (c, c -> b)) -> Diff p b c
forall a b. (a -> b) -> a -> b
$ \b
b ->
  let dim :: Int
dim = [j] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [j]
x0
      -- Forward: iterate from caller-supplied seed
      stepFwd :: [j] -> [j]
stepFwd [j]
x = let (([j]
x', c
_), ([j], c) -> ([j], b)
_) = ([j], b) -> (([j], c), ([j], c) -> ([j], b))
body ([j]
x, b
b) in [j]
x'
      a :: [j]
a = ([j] -> [j]) -> [j] -> [[j]]
forall a. (a -> a) -> a -> [a]
iterate [j] -> [j]
stepFwd [j]
x0 [[j]] -> Int -> [j]
forall a. HasCallStack => [a] -> Int -> a
!! Int
n
      (([j]
_, c
c), ([j], c) -> ([j], b)
backward) = ([j], b) -> (([j], c), ([j], c) -> ([j], b))
body ([j]
a, b
b)
      -- Channel self-coupling, column i = backward probe at e_i
      zeroV :: [j]
zeroV = Int -> j -> [j]
forall a. Int -> a -> [a]
replicate Int
dim j
forall a. Additive a => a
NHA.zero
      basisVec :: Int -> [j]
basisVec Int
i = [if Int
k Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
i then j
forall a. Multiplicative a => a
NHM.one else j
forall a. Additive a => a
NHA.zero | Int
k <- [Int
0 .. Int
dim Int -> Int -> Int
forall a. Subtractive a => a -> a -> a
- Int
1]]
      cols :: [[j]]
cols = [([j], b) -> [j]
forall a b. (a, b) -> a
fst (([j], c) -> ([j], b)
backward (Int -> [j]
basisVec Int
i, () -> c
forall {k} (arr :: * -> k -> *) (a :: k). Zero arr a => arr () a
CB.zero ())) | Int
i <- [Int
0 .. Int
dim Int -> Int -> Int
forall a. Subtractive a => a -> a -> a
- Int
1]]
      aMat :: Matrix j
aMat = [[j]] -> Matrix j
forall a. [[a]] -> Matrix a
fromLists [[[j]
col [j] -> Int -> j
forall a. HasCallStack => [a] -> Int -> a
!! Int
k | [j]
col <- [[j]]
cols] | Int
k <- [Int
0 .. Int
dim Int -> Int -> Int
forall a. Subtractive a => a -> a -> a
- Int
1]]
      -- star A — Gaussian elimination / Warshall / Floyd–Warshall /
      -- state elimination, depending on the carrier
      aStar :: Matrix j
aStar = Matrix j -> Matrix j
forall a. StarSemiring a => Matrix a -> Matrix a
MD.starMatrix Matrix j
aMat
      pullback :: c -> b
pullback c
dc =
        let cdc :: [j]
cdc = ([j], b) -> [j]
forall a b. (a, b) -> a
fst (([j], c) -> ([j], b)
backward ([j]
zeroV, c
dc))
         in ([j], b) -> b
forall a b. (a, b) -> b
snd (([j], c) -> ([j], b)
backward (Matrix j -> [j] -> [j]
forall a. (Additive a, Multiplicative a) => Matrix a -> [a] -> [a]
matVec Matrix j
aStar [j]
cdc, c
dc))
   in (c
c, c -> b
pullback)

-- | 'Circuit.Diff.Circuit.traceStarFrom' for a bare 'Double' channel, routed
-- through 'FieldStar' so @star a = recip (1 − a)@ is the class
-- method, not an ad-hoc formula.
--
-- Channel @x' = 0.3x + b@, output @c = 2x@.  The fixpoint is
-- @x = b\/0.7@, so @dc\/db = 2\/0.7 ≈ 2.857@ — the same number the
-- Neumann iteration approaches, computed here in closed form:
--
-- >>> :{
-- let body = Diff (\(x, b) ->
--       ( (0.3 * x + b, 2.0 * x)
--       , \(dx', dc) -> (0.3 * dx' + 2.0 * dc, dx') ))
-- :}
--
-- >>> let (y, pb) = runDiff (traceStarFromD 0 60 body) 1.0
-- >>> abs (y - 2.0 / 0.7) < 1e-12
-- True
-- >>> abs (pb 1.0 - 2.0 / 0.7) < 1e-12
-- True
traceStarFromD ::
  (MergeZero (->) c) =>
  Double ->
  Int ->
  Diff p (Double, b) (Double, c) ->
  Diff p b c
traceStarFromD :: forall {k} c (p :: k) b.
MergeZero (->) c =>
Double -> Int -> Diff p (Double, b) (Double, c) -> Diff p b c
traceStarFromD Double
x0 Int
n (Diff (Double, b) -> ((Double, c), (Double, c) -> (Double, b))
body) =
  FieldStar
-> Int -> Diff p (FieldStar, b) (FieldStar, c) -> Diff p b c
forall {k} j c (p :: k) b.
(StarSemiring j, MergeZero (->) c) =>
j -> Int -> Diff p (j, b) (j, c) -> Diff p b c
traceStarFrom (Double -> FieldStar
FieldStar Double
x0) Int
n (((FieldStar, b)
 -> ((FieldStar, c), (FieldStar, c) -> (FieldStar, b)))
-> Diff p (FieldStar, b) (FieldStar, c)
forall k (p :: k) a b. (a -> (b, b -> a)) -> Diff p a b
Diff (FieldStar, b)
-> ((FieldStar, c), (FieldStar, c) -> (FieldStar, b))
body')
  where
    body' :: (FieldStar, b)
-> ((FieldStar, c), (FieldStar, c) -> (FieldStar, b))
body' (FieldStar Double
x, b
b) =
      let ((Double
x', c
c), (Double, c) -> (Double, b)
back) = (Double, b) -> ((Double, c), (Double, c) -> (Double, b))
body (Double
x, b
b)
       in ( (Double -> FieldStar
FieldStar Double
x', c
c),
            \(FieldStar Double
dx', c
dc) ->
              let (Double
dx, b
db) = (Double, c) -> (Double, b)
back (Double
dx', c
dc)
               in (Double -> FieldStar
FieldStar Double
dx, b
db)
          )

-- | 'traceStarMatrix' for bare @[Double]@ channels, routed through
-- 'FieldStar' so 'starMatrix' performs honest @(I − A)⁻¹@.
--
-- A two-dimensional channel: @x' = 0.5y + b@, @y' = 0.3x@, output
-- @c = x + 2y@.  Solving the fixpoint by hand: @x = b\/0.85@,
-- @y = 0.3b\/0.85@, @c = 1.6b\/0.85@, so @dc\/db = 1.6\/0.85@.
--
-- >>> :{
-- let body2 = Diff (\([x, y], b) ->
--       ( ([0.5 * y + b, 0.3 * x], x + 2.0 * y)
--       , \([dx', dy'], dc) -> ([0.3 * dy' + dc, 0.5 * dx' + 2.0 * dc], dx') ))
-- :}
--
-- >>> let (y2, pb2) = runDiff (traceStarMatrixD [0, 0] 200 body2) 1.0
-- >>> abs (y2 - 1.6 / 0.85) < 1e-12
-- True
-- >>> abs (pb2 1.0 - 1.6 / 0.85) < 1e-12
-- True
traceStarMatrixD ::
  (MergeZero (->) c) =>
  [Double] ->
  Int ->
  Diff p ([Double], b) ([Double], c) ->
  Diff p b c
traceStarMatrixD :: forall {k} c (p :: k) b.
MergeZero (->) c =>
[Double] -> Int -> Diff p ([Double], b) ([Double], c) -> Diff p b c
traceStarMatrixD [Double]
x0 Int
n (Diff ([Double], b) -> (([Double], c), ([Double], c) -> ([Double], b))
body) =
  [FieldStar]
-> Int -> Diff p ([FieldStar], b) ([FieldStar], c) -> Diff p b c
forall {k} j c (p :: k) b.
(StarSemiring j, MergeZero (->) c) =>
[j] -> Int -> Diff p ([j], b) ([j], c) -> Diff p b c
traceStarMatrix ((Double -> FieldStar) -> [Double] -> [FieldStar]
forall a b. (a -> b) -> [a] -> [b]
map Double -> FieldStar
FieldStar [Double]
x0) Int
n ((([FieldStar], b)
 -> (([FieldStar], c), ([FieldStar], c) -> ([FieldStar], b)))
-> Diff p ([FieldStar], b) ([FieldStar], c)
forall k (p :: k) a b. (a -> (b, b -> a)) -> Diff p a b
Diff ([FieldStar], b)
-> (([FieldStar], c), ([FieldStar], c) -> ([FieldStar], b))
body')
  where
    wrap :: [Double] -> [FieldStar]
wrap = (Double -> FieldStar) -> [Double] -> [FieldStar]
forall a b. (a -> b) -> [a] -> [b]
map Double -> FieldStar
FieldStar
    unwrap :: [FieldStar] -> [Double]
unwrap = (FieldStar -> Double) -> [FieldStar] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (\(FieldStar Double
d) -> Double
d)
    body' :: ([FieldStar], b)
-> (([FieldStar], c), ([FieldStar], c) -> ([FieldStar], b))
body' ([FieldStar]
js, b
b) =
      let (([Double]
xs, c
c), ([Double], c) -> ([Double], b)
back) = ([Double], b) -> (([Double], c), ([Double], c) -> ([Double], b))
body ([FieldStar] -> [Double]
unwrap [FieldStar]
js, b
b)
       in ( ([Double] -> [FieldStar]
wrap [Double]
xs, c
c),
            \([FieldStar]
djs, c
dc) ->
              let ([Double]
dxs, b
db) = ([Double], c) -> ([Double], b)
back ([FieldStar] -> [Double]
unwrap [FieldStar]
djs, c
dc)
               in ([Double] -> [FieldStar]
wrap [Double]
dxs, b
db)
          )

-- | 'FieldStar' as a numeric carrier for circuits' additive structure, so
-- hand-built bodies can use structural rows at 'FieldStar' types.  (Bodies from
-- 'Circuit.Diff.Backprop.linearizeBody' never need this — their structural rows
-- are already pointwise pullbacks.)
--
-- These instances remain necessary because structural rows ('Plus', 'Zero')
-- carry 'MergeT' / 'ZeroT' constraints that resolve to 'Merge' / 'Zero' on
-- the wiring arrow.  For 'Pullback FieldStar' those constraints bottom out in
-- 'Merge (->) FieldStar' and 'Zero (->) FieldStar'.  Deliberately orphan: this
-- module is the federation seam between @circuits@ and @numhask-free@.
instance Circuit.Bimonoid.Merge (->) FieldStar where
  plus :: (FieldStar, FieldStar) -> FieldStar
plus (FieldStar Double
x, FieldStar Double
y) = Double -> FieldStar
FieldStar (Double
x Double -> Double -> Double
forall a. Additive a => a -> a -> a
+ Double
y)

instance Circuit.Bimonoid.Zero (->) FieldStar where
  zero :: () -> FieldStar
zero ()
_ = Double -> FieldStar
FieldStar Double
0

-- | Solve one affine body in closed form.
--
-- For a body @f :: (StarChannel j, c) -> (StarChannel j, b)@ (channel
-- cotangent, output cotangent), affinity gives @f₁ (dj, dc) = A·dj + C·dc@.
-- The probes:
--
-- > C·dc = f₁ (0, dc)                       -- one call, at the true zero
-- > A·e_i = f₁ (e_i, dc) − C·dc             -- offset-subtracted: valid at dc ≠ 0
-- > dj   = star A · C·dc                    -- starMatrix
-- > db   = f₂ (dj, dc)                      -- one final call
--
-- The dictionary is recovered from the output of a lazy probe: a body built
-- with 'Circuit.Diff.Evidence.withStarChannel' carries the same dictionary
-- through the trace, so pattern-matching the output reveals it without
-- evaluating the input data.
--
-- The offset subtraction is what lets every probe run at the /actual/
-- cotangent @dc@, so no @zero@ for the (existential) type @c@ is ever
-- needed.  Cost: @dim + 2@ body calls per cotangent; the star is /not/
-- shared across cotangents — that is the price of the existential.
solveAffine ::
  forall j c b.
  ((StarChannel j, c) -> (StarChannel j, b)) ->
  c ->
  b
solveAffine :: forall j c b. ((StarChannel j, c) -> (StarChannel j, b)) -> c -> b
solveAffine (StarChannel j, c) -> (StarChannel j, b)
body c
dc =
  let -- A bottom value of type 'StarChannel j'.  Only the constructor is
      -- needed; the fields are supplied lazily by the body's output.
      probe :: StarChannel j
      probe :: StarChannel j
probe =
        StarChannel
          { starDim :: Int
starDim = [Char] -> Int
forall a. HasCallStack => [Char] -> a
error [Char]
"Circuit.Diff.Star.solveAffine: probe dim evaluated",
            starData :: j
starData = [Char] -> j
forall a. HasCallStack => [Char] -> a
error [Char]
"Circuit.Diff.Star.solveAffine: probe data evaluated",
            starZero :: Int -> j
starZero = [Char] -> Int -> j
forall a. HasCallStack => [Char] -> a
error [Char]
"Circuit.Diff.Star.solveAffine: probe zero evaluated",
            starBasis :: Int -> Int -> j
starBasis = [Char] -> Int -> Int -> j
forall a. HasCallStack => [Char] -> a
error [Char]
"Circuit.Diff.Star.solveAffine: probe basis evaluated",
            starAdd :: j -> j -> j
starAdd = [Char] -> j -> j -> j
forall a. HasCallStack => [Char] -> a
error [Char]
"Circuit.Diff.Star.solveAffine: probe add evaluated",
            starNegate :: j -> j
starNegate = [Char] -> j -> j
forall a. HasCallStack => [Char] -> a
error [Char]
"Circuit.Diff.Star.solveAffine: probe negate evaluated",
            starSelfMatrix :: Int -> (j -> j) -> Matrix (Scalar j)
starSelfMatrix = [Char] -> Int -> (j -> j) -> Matrix (Scalar j)
forall a. HasCallStack => [Char] -> a
error [Char]
"Circuit.Diff.Star.solveAffine: probe selfMatrix evaluated",
            starApplyMatrix :: Matrix (Scalar j) -> j -> j
starApplyMatrix = [Char] -> Matrix (Scalar j) -> j -> j
forall a. HasCallStack => [Char] -> a
error [Char]
"Circuit.Diff.Star.solveAffine: probe applyMatrix evaluated",
            starMatrix :: Matrix (Scalar j) -> Matrix (Scalar j)
starMatrix = [Char] -> Matrix (Scalar j) -> Matrix (Scalar j)
forall a. HasCallStack => [Char] -> a
error [Char]
"Circuit.Diff.Star.solveAffine: probe matrix evaluated"
          }
      (StarChannel j
scOut, b
_) = (StarChannel j, c) -> (StarChannel j, b)
body (StarChannel j
probe, c
dc)
      zeroJ :: j
zeroJ = StarChannel j -> Int -> j
forall s. StarChannel s -> Int -> s
starZero StarChannel j
scOut (StarChannel j -> Int
forall s. StarChannel s -> Int
starDim StarChannel j
scOut)
      zeroSC :: StarChannel j
zeroSC = StarChannel j
scOut {starData = zeroJ}
      cdc :: j
cdc = StarChannel j -> j
forall s. StarChannel s -> s
starData ((StarChannel j, b) -> StarChannel j
forall a b. (a, b) -> a
fst ((StarChannel j, c) -> (StarChannel j, b)
body (StarChannel j
zeroSC, c
dc)))
      negCdc :: j
negCdc = StarChannel j -> j -> j
forall s. StarChannel s -> s -> s
starNegate StarChannel j
scOut j
cdc
      aMat :: Matrix (Scalar j)
aMat =
        StarChannel j -> Int -> (j -> j) -> Matrix (Scalar j)
forall s. StarChannel s -> Int -> (s -> s) -> Matrix (Scalar s)
starSelfMatrix
          StarChannel j
scOut
          (StarChannel j -> Int
forall s. StarChannel s -> Int
starDim StarChannel j
scOut)
          (\j
dk -> StarChannel j -> j -> j -> j
forall s. StarChannel s -> s -> s -> s
starAdd StarChannel j
scOut (StarChannel j -> j
forall s. StarChannel s -> s
starData ((StarChannel j, b) -> StarChannel j
forall a b. (a, b) -> a
fst ((StarChannel j, c) -> (StarChannel j, b)
body (StarChannel j
zeroSC {starData = dk}, c
dc)))) j
negCdc)
      dj :: j
dj = StarChannel j -> Matrix (Scalar j) -> j -> j
forall s. StarChannel s -> Matrix (Scalar s) -> s -> s
starApplyMatrix StarChannel j
scOut (StarChannel j -> Matrix (Scalar j) -> Matrix (Scalar j)
forall s. StarChannel s -> Matrix (Scalar s) -> Matrix (Scalar s)
starMatrix StarChannel j
scOut Matrix (Scalar j)
aMat) j
cdc
   in (StarChannel j, b) -> b
forall a b. (a, b) -> b
snd ((StarChannel j, c) -> (StarChannel j, b)
body (StarChannel j
zeroSC {starData = dj}, c
dc))

-- | Solve a @(,)@ feedback body whose channel is 'StarChannel' in closed form.
--
-- The channel type stays exposed, so the dictionary is read directly from the
-- 'StarChannel' value carried on the feedback wire.
--
-- A self-coupled scalar body the lazy trace diverges on, solved exactly
-- (@dj = 0.3·dj + 2·dc@, @db = dj@, so @db\/dc = 2\/0.7@):
--
-- >>> :{
-- let body (FieldStar dj, dc) = (FieldStar (0.3 * dj + 2.0 * dc), dj)
--     b = Body (withStarChannel fieldStarChannel (Pullback body)) :: Body (,) (StarChannel FieldStar) Pullback Double Double
-- :}
--
-- >>> let solved = solveStarBody b
-- >>> abs (runPullback solved 1.0 - 2.0 / 0.7) < 1e-12
-- True
solveStarBody ::
  forall s a b.
  Body (,) (StarChannel s) Pullback b a ->
  Pullback b a
solveStarBody :: forall s a b. Body (,) (StarChannel s) Pullback b a -> Pullback b a
solveStarBody (Body Pullback (StarChannel s, b) (StarChannel s, a)
f) = (b -> a) -> Pullback b a
forall b a. (b -> a) -> Pullback b a
Pullback (((StarChannel s, b) -> (StarChannel s, a)) -> b -> a
forall j c b. ((StarChannel j, c) -> (StarChannel j, b)) -> c -> b
solveAffine (Pullback (StarChannel s, b) (StarChannel s, a)
-> (StarChannel s, b) -> (StarChannel s, a)
forall b a. Pullback b a -> b -> a
runPullback Pullback (StarChannel s, b) (StarChannel s, a)
f))