-- | Seat-level tensor combinators over 'Ends' with state carried by
-- 'Body'.
--
-- These are the semantic citizens that free-agent 'FreeSeat' terms fold into.
-- The buffer state is part of the base arrow ('Body (,) (K IO)
-- [Post Text]') rather than a monad transformer.
module Circuit.Agent.Tensor
  ( AgentShard,
    silentShard,
    awaitShard,
    raceShard,
    fanOutShard,
    fanInShard,
    runSubShard,
    closeShardIO,
    ioShard,
    writeBatch,
    readBatch,
    synthesisSummary,
  )
where

import Circuit (Body (..), close, companion, conjoint)
import Circuit.Agent (Name, Post (..), PostId, mkPost, synthesis)
import Circuit.Category (K (..))
import Circuit.Poles (Poles, poles0)
import Data.Text (Text)
import Data.Text qualified as T

-- | An effectful agent shard with a '[Post Text]' buffer threaded through the
-- arrow.
type AgentShard a b = Poles (Body (,) [Post Text] (K IO)) a b

-- | Run a child shard in isolation on the given input, discarding its
-- residual state.  Branches in await / race / fan-out have private scratch
-- state; only their emits rejoin the main stream.
runSubShard :: AgentShard [Post Text] [Post Text] -> [Post Text] -> IO [Post Text]
runSubShard :: AgentShard [Post Text] [Post Text] -> [Post Text] -> IO [Post Text]
runSubShard AgentShard [Post Text] [Post Text]
sh [Post Text]
xs = do
  let thread :: Body (,) [Post Text] (K IO) [Post Text] [Post Text]
thread = In (Body (,) [Post Text] (K IO)) [Post Text]
-> Out (Body (,) [Post Text] (K IO)) [Post Text]
-> Body (,) [Post Text] (K IO) [Post Text] [Post Text]
forall {k} (arr :: k -> k -> *) (a :: k).
In arr a -> Out arr a -> arr a a
close (AgentShard [Post Text] [Post Text]
-> In (Body (,) [Post Text] (K IO)) [Post Text]
forall {k1} {k2} (arr :: k1 -> k2 -> *) (a :: k1) (b :: k2).
Poles arr a b -> In arr a
conjoint AgentShard [Post Text] [Post Text]
sh) (AgentShard [Post Text] [Post Text]
-> Out (Body (,) [Post Text] (K IO)) [Post Text]
forall {k1} {k2} (arr :: k1 -> k2 -> *) (a :: k1) (b :: k2).
Poles arr a b -> Out arr b
companion AgentShard [Post Text] [Post Text]
sh)
  (_, ys) <- K IO ([Post Text], [Post Text]) ([Post Text], [Post Text])
-> ([Post Text], [Post Text]) -> IO ([Post Text], [Post Text])
forall {k} (m :: k -> *) a (b :: k). K m a b -> a -> m b
runK (Body (,) [Post Text] (K IO) [Post Text] [Post Text]
-> K IO ([Post Text], [Post Text]) ([Post Text], [Post Text])
forall {k1} {k2} {k3} (t :: k1 -> k2 -> k3) (ch :: k1)
       (arr :: k3 -> k3 -> *) (a :: k2) (b :: k2).
Body t ch arr a b -> arr (t ch a) (t ch b)
morphism Body (,) [Post Text] (K IO) [Post Text] [Post Text]
thread) ([], [Post Text]
xs)
  pure ys

-- | Close a same-type shard once, exposing both the output and the residual
-- state.  Works for any state carrier @s@ and payload @a@.
closeShardIO :: Poles (Body (,) s (K IO)) a a -> a -> s -> IO (a, s)
closeShardIO :: forall s a. Poles (Body (,) s (K IO)) a a -> a -> s -> IO (a, s)
closeShardIO Poles (Body (,) s (K IO)) a a
sh a
x s
s0 = do
  let thread :: Body (,) s (K IO) a a
thread = In (Body (,) s (K IO)) a
-> Out (Body (,) s (K IO)) a -> Body (,) s (K IO) a a
forall {k} (arr :: k -> k -> *) (a :: k).
In arr a -> Out arr a -> arr a a
close (Poles (Body (,) s (K IO)) a a -> In (Body (,) s (K IO)) a
forall {k1} {k2} (arr :: k1 -> k2 -> *) (a :: k1) (b :: k2).
Poles arr a b -> In arr a
conjoint Poles (Body (,) s (K IO)) a a
sh) (Poles (Body (,) s (K IO)) a a -> Out (Body (,) s (K IO)) a
forall {k1} {k2} (arr :: k1 -> k2 -> *) (a :: k1) (b :: k2).
Poles arr a b -> Out arr b
companion Poles (Body (,) s (K IO)) a a
sh)
  (s', y) <- K IO (s, a) (s, a) -> (s, a) -> IO (s, a)
forall {k} (m :: k -> *) a (b :: k). K m a b -> a -> m b
runK (Body (,) s (K IO) a a -> K IO (s, a) (s, a)
forall {k1} {k2} {k3} (t :: k1 -> k2 -> k3) (ch :: k1)
       (arr :: k3 -> k3 -> *) (a :: k2) (b :: k2).
Body t ch arr a b -> arr (t ch a) (t ch b)
morphism Body (,) s (K IO) a a
thread) (s
s0, a
x)
  pure (y, s')

-- | Helper: store the input batch as the buffer.
writeBatch :: Body (,) [Post Text] (K IO) [Post Text] ()
writeBatch :: Body (,) [Post Text] (K IO) [Post Text] ()
writeBatch = K IO ([Post Text], [Post Text]) ([Post Text], ())
-> Body (,) [Post Text] (K IO) [Post Text] ()
forall {k} {k1} {k2} (t :: k -> k1 -> k2) (ch :: k)
       (arr :: k2 -> k2 -> *) (a :: k1) (b :: k1).
arr (t ch a) (t ch b) -> Body t ch arr a b
Body (K IO ([Post Text], [Post Text]) ([Post Text], ())
 -> Body (,) [Post Text] (K IO) [Post Text] ())
-> K IO ([Post Text], [Post Text]) ([Post Text], ())
-> Body (,) [Post Text] (K IO) [Post Text] ()
forall a b. (a -> b) -> a -> b
$ (([Post Text], [Post Text]) -> IO ([Post Text], ()))
-> K IO ([Post Text], [Post Text]) ([Post Text], ())
forall {k} (m :: k -> *) a (b :: k). (a -> m b) -> K m a b
K ((([Post Text], [Post Text]) -> IO ([Post Text], ()))
 -> K IO ([Post Text], [Post Text]) ([Post Text], ()))
-> (([Post Text], [Post Text]) -> IO ([Post Text], ()))
-> K IO ([Post Text], [Post Text]) ([Post Text], ())
forall a b. (a -> b) -> a -> b
$ \([Post Text]
_, [Post Text]
xs) -> ([Post Text], ()) -> IO ([Post Text], ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Post Text]
xs, ())

-- | Helper: return the buffer as output and clear it.
readBatch :: Body (,) [Post Text] (K IO) () [Post Text]
readBatch :: Body (,) [Post Text] (K IO) () [Post Text]
readBatch = K IO ([Post Text], ()) ([Post Text], [Post Text])
-> Body (,) [Post Text] (K IO) () [Post Text]
forall {k} {k1} {k2} (t :: k -> k1 -> k2) (ch :: k)
       (arr :: k2 -> k2 -> *) (a :: k1) (b :: k1).
arr (t ch a) (t ch b) -> Body t ch arr a b
Body (K IO ([Post Text], ()) ([Post Text], [Post Text])
 -> Body (,) [Post Text] (K IO) () [Post Text])
-> K IO ([Post Text], ()) ([Post Text], [Post Text])
-> Body (,) [Post Text] (K IO) () [Post Text]
forall a b. (a -> b) -> a -> b
$ (([Post Text], ()) -> IO ([Post Text], [Post Text]))
-> K IO ([Post Text], ()) ([Post Text], [Post Text])
forall {k} (m :: k -> *) a (b :: k). (a -> m b) -> K m a b
K ((([Post Text], ()) -> IO ([Post Text], [Post Text]))
 -> K IO ([Post Text], ()) ([Post Text], [Post Text]))
-> (([Post Text], ()) -> IO ([Post Text], [Post Text]))
-> K IO ([Post Text], ()) ([Post Text], [Post Text])
forall a b. (a -> b) -> a -> b
$ \([Post Text]
s, ()) -> ([Post Text], [Post Text]) -> IO ([Post Text], [Post Text])
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([], [Post Text]
s)

-- | Build an agent shard that stores the input batch, then computes outputs
-- from that batch in 'IO'.
ioShard :: ([Post Text] -> IO [Post Text]) -> AgentShard [Post Text] [Post Text]
ioShard :: ([Post Text] -> IO [Post Text])
-> AgentShard [Post Text] [Post Text]
ioShard [Post Text] -> IO [Post Text]
emit = Body (,) [Post Text] (K IO) [Post Text] ()
-> Body (,) [Post Text] (K IO) () [Post Text]
-> AgentShard [Post Text] [Post Text]
forall (arr :: * -> * -> *) a b.
HasDual () arr =>
arr a () -> arr () b -> Poles arr a b
poles0 Body (,) [Post Text] (K IO) [Post Text] ()
writeBatch Body (,) [Post Text] (K IO) () [Post Text]
read'
  where
    read' :: Body (,) [Post Text] (K IO) () [Post Text]
read' = K IO ([Post Text], ()) ([Post Text], [Post Text])
-> Body (,) [Post Text] (K IO) () [Post Text]
forall {k} {k1} {k2} (t :: k -> k1 -> k2) (ch :: k)
       (arr :: k2 -> k2 -> *) (a :: k1) (b :: k1).
arr (t ch a) (t ch b) -> Body t ch arr a b
Body (K IO ([Post Text], ()) ([Post Text], [Post Text])
 -> Body (,) [Post Text] (K IO) () [Post Text])
-> K IO ([Post Text], ()) ([Post Text], [Post Text])
-> Body (,) [Post Text] (K IO) () [Post Text]
forall a b. (a -> b) -> a -> b
$ (([Post Text], ()) -> IO ([Post Text], [Post Text]))
-> K IO ([Post Text], ()) ([Post Text], [Post Text])
forall {k} (m :: k -> *) a (b :: k). (a -> m b) -> K m a b
K ((([Post Text], ()) -> IO ([Post Text], [Post Text]))
 -> K IO ([Post Text], ()) ([Post Text], [Post Text]))
-> (([Post Text], ()) -> IO ([Post Text], [Post Text]))
-> K IO ([Post Text], ()) ([Post Text], [Post Text])
forall a b. (a -> b) -> a -> b
$ \([Post Text]
s, ()) -> do
      outs <- [Post Text] -> IO [Post Text]
emit [Post Text]
s
      pure ([], outs)

-- | Silent shard: commit replaces the buffer; emit clears it and returns [].
silentShard :: AgentShard [Post Text] [Post Text]
silentShard :: AgentShard [Post Text] [Post Text]
silentShard = ([Post Text] -> IO [Post Text])
-> AgentShard [Post Text] [Post Text]
ioShard ([Post Text] -> IO [Post Text]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Post Text] -> IO [Post Text])
-> ([Post Text] -> [Post Text]) -> [Post Text] -> IO [Post Text]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Post Text] -> [Post Text] -> [Post Text]
forall a b. a -> b -> a
const [])

-- | Product / await shard: both sub-shards see the same input; emits are
-- concatenated left-to-right.
awaitShard ::
  AgentShard [Post Text] [Post Text] ->
  AgentShard [Post Text] [Post Text] ->
  AgentShard [Post Text] [Post Text]
awaitShard :: AgentShard [Post Text] [Post Text]
-> AgentShard [Post Text] [Post Text]
-> AgentShard [Post Text] [Post Text]
awaitShard AgentShard [Post Text] [Post Text]
sh1 AgentShard [Post Text] [Post Text]
sh2 = Body (,) [Post Text] (K IO) [Post Text] ()
-> Body (,) [Post Text] (K IO) () [Post Text]
-> AgentShard [Post Text] [Post Text]
forall (arr :: * -> * -> *) a b.
HasDual () arr =>
arr a () -> arr () b -> Poles arr a b
poles0 Body (,) [Post Text] (K IO) [Post Text] ()
writeBatch Body (,) [Post Text] (K IO) () [Post Text]
read'
  where
    read' :: Body (,) [Post Text] (K IO) () [Post Text]
read' = K IO ([Post Text], ()) ([Post Text], [Post Text])
-> Body (,) [Post Text] (K IO) () [Post Text]
forall {k} {k1} {k2} (t :: k -> k1 -> k2) (ch :: k)
       (arr :: k2 -> k2 -> *) (a :: k1) (b :: k1).
arr (t ch a) (t ch b) -> Body t ch arr a b
Body (K IO ([Post Text], ()) ([Post Text], [Post Text])
 -> Body (,) [Post Text] (K IO) () [Post Text])
-> K IO ([Post Text], ()) ([Post Text], [Post Text])
-> Body (,) [Post Text] (K IO) () [Post Text]
forall a b. (a -> b) -> a -> b
$ (([Post Text], ()) -> IO ([Post Text], [Post Text]))
-> K IO ([Post Text], ()) ([Post Text], [Post Text])
forall {k} (m :: k -> *) a (b :: k). (a -> m b) -> K m a b
K ((([Post Text], ()) -> IO ([Post Text], [Post Text]))
 -> K IO ([Post Text], ()) ([Post Text], [Post Text]))
-> (([Post Text], ()) -> IO ([Post Text], [Post Text]))
-> K IO ([Post Text], ()) ([Post Text], [Post Text])
forall a b. (a -> b) -> a -> b
$ \([Post Text]
s, ()) -> do
      o1 <- AgentShard [Post Text] [Post Text] -> [Post Text] -> IO [Post Text]
runSubShard AgentShard [Post Text] [Post Text]
sh1 [Post Text]
s
      o2 <- runSubShard sh2 s
      pure ([], o1 ++ o2)

-- | Coproduct / race shard: left emit wins if non-empty, otherwise right.
raceShard ::
  AgentShard [Post Text] [Post Text] ->
  AgentShard [Post Text] [Post Text] ->
  AgentShard [Post Text] [Post Text]
raceShard :: AgentShard [Post Text] [Post Text]
-> AgentShard [Post Text] [Post Text]
-> AgentShard [Post Text] [Post Text]
raceShard AgentShard [Post Text] [Post Text]
sh1 AgentShard [Post Text] [Post Text]
sh2 = Body (,) [Post Text] (K IO) [Post Text] ()
-> Body (,) [Post Text] (K IO) () [Post Text]
-> AgentShard [Post Text] [Post Text]
forall (arr :: * -> * -> *) a b.
HasDual () arr =>
arr a () -> arr () b -> Poles arr a b
poles0 Body (,) [Post Text] (K IO) [Post Text] ()
writeBatch Body (,) [Post Text] (K IO) () [Post Text]
read'
  where
    read' :: Body (,) [Post Text] (K IO) () [Post Text]
read' = K IO ([Post Text], ()) ([Post Text], [Post Text])
-> Body (,) [Post Text] (K IO) () [Post Text]
forall {k} {k1} {k2} (t :: k -> k1 -> k2) (ch :: k)
       (arr :: k2 -> k2 -> *) (a :: k1) (b :: k1).
arr (t ch a) (t ch b) -> Body t ch arr a b
Body (K IO ([Post Text], ()) ([Post Text], [Post Text])
 -> Body (,) [Post Text] (K IO) () [Post Text])
-> K IO ([Post Text], ()) ([Post Text], [Post Text])
-> Body (,) [Post Text] (K IO) () [Post Text]
forall a b. (a -> b) -> a -> b
$ (([Post Text], ()) -> IO ([Post Text], [Post Text]))
-> K IO ([Post Text], ()) ([Post Text], [Post Text])
forall {k} (m :: k -> *) a (b :: k). (a -> m b) -> K m a b
K ((([Post Text], ()) -> IO ([Post Text], [Post Text]))
 -> K IO ([Post Text], ()) ([Post Text], [Post Text]))
-> (([Post Text], ()) -> IO ([Post Text], [Post Text]))
-> K IO ([Post Text], ()) ([Post Text], [Post Text])
forall a b. (a -> b) -> a -> b
$ \([Post Text]
s, ()) -> do
      o1 <- AgentShard [Post Text] [Post Text] -> [Post Text] -> IO [Post Text]
runSubShard AgentShard [Post Text] [Post Text]
sh1 [Post Text]
s
      o2 <- runSubShard sh2 s
      pure ([], if null o1 then o2 else o1)

-- | Fan-out shard: every sub-shard sees the same input; emits are concatenated
-- in branch order.
fanOutShard ::
  [AgentShard [Post Text] [Post Text]] ->
  AgentShard [Post Text] [Post Text]
fanOutShard :: [AgentShard [Post Text] [Post Text]]
-> AgentShard [Post Text] [Post Text]
fanOutShard [AgentShard [Post Text] [Post Text]]
shs = Body (,) [Post Text] (K IO) [Post Text] ()
-> Body (,) [Post Text] (K IO) () [Post Text]
-> AgentShard [Post Text] [Post Text]
forall (arr :: * -> * -> *) a b.
HasDual () arr =>
arr a () -> arr () b -> Poles arr a b
poles0 Body (,) [Post Text] (K IO) [Post Text] ()
writeBatch Body (,) [Post Text] (K IO) () [Post Text]
read'
  where
    read' :: Body (,) [Post Text] (K IO) () [Post Text]
read' = K IO ([Post Text], ()) ([Post Text], [Post Text])
-> Body (,) [Post Text] (K IO) () [Post Text]
forall {k} {k1} {k2} (t :: k -> k1 -> k2) (ch :: k)
       (arr :: k2 -> k2 -> *) (a :: k1) (b :: k1).
arr (t ch a) (t ch b) -> Body t ch arr a b
Body (K IO ([Post Text], ()) ([Post Text], [Post Text])
 -> Body (,) [Post Text] (K IO) () [Post Text])
-> K IO ([Post Text], ()) ([Post Text], [Post Text])
-> Body (,) [Post Text] (K IO) () [Post Text]
forall a b. (a -> b) -> a -> b
$ (([Post Text], ()) -> IO ([Post Text], [Post Text]))
-> K IO ([Post Text], ()) ([Post Text], [Post Text])
forall {k} (m :: k -> *) a (b :: k). (a -> m b) -> K m a b
K ((([Post Text], ()) -> IO ([Post Text], [Post Text]))
 -> K IO ([Post Text], ()) ([Post Text], [Post Text]))
-> (([Post Text], ()) -> IO ([Post Text], [Post Text]))
-> K IO ([Post Text], ()) ([Post Text], [Post Text])
forall a b. (a -> b) -> a -> b
$ \([Post Text]
s, ()) -> do
      os <- (AgentShard [Post Text] [Post Text] -> IO [Post Text])
-> [AgentShard [Post Text] [Post Text]] -> IO [[Post Text]]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse (AgentShard [Post Text] [Post Text] -> [Post Text] -> IO [Post Text]
`runSubShard` [Post Text]
s) [AgentShard [Post Text] [Post Text]]
shs
      pure ([], concat os)

-- | Fan-in shard: fan-out, then collapse the collected branch outputs with the
-- supplied summary function.
fanInShard ::
  ([[Post Text]] -> [Post Text]) ->
  [AgentShard [Post Text] [Post Text]] ->
  AgentShard [Post Text] [Post Text]
fanInShard :: ([[Post Text]] -> [Post Text])
-> [AgentShard [Post Text] [Post Text]]
-> AgentShard [Post Text] [Post Text]
fanInShard [[Post Text]] -> [Post Text]
summary [AgentShard [Post Text] [Post Text]]
shs = Body (,) [Post Text] (K IO) [Post Text] ()
-> Body (,) [Post Text] (K IO) () [Post Text]
-> AgentShard [Post Text] [Post Text]
forall (arr :: * -> * -> *) a b.
HasDual () arr =>
arr a () -> arr () b -> Poles arr a b
poles0 Body (,) [Post Text] (K IO) [Post Text] ()
writeBatch Body (,) [Post Text] (K IO) () [Post Text]
read'
  where
    read' :: Body (,) [Post Text] (K IO) () [Post Text]
read' = K IO ([Post Text], ()) ([Post Text], [Post Text])
-> Body (,) [Post Text] (K IO) () [Post Text]
forall {k} {k1} {k2} (t :: k -> k1 -> k2) (ch :: k)
       (arr :: k2 -> k2 -> *) (a :: k1) (b :: k1).
arr (t ch a) (t ch b) -> Body t ch arr a b
Body (K IO ([Post Text], ()) ([Post Text], [Post Text])
 -> Body (,) [Post Text] (K IO) () [Post Text])
-> K IO ([Post Text], ()) ([Post Text], [Post Text])
-> Body (,) [Post Text] (K IO) () [Post Text]
forall a b. (a -> b) -> a -> b
$ (([Post Text], ()) -> IO ([Post Text], [Post Text]))
-> K IO ([Post Text], ()) ([Post Text], [Post Text])
forall {k} (m :: k -> *) a (b :: k). (a -> m b) -> K m a b
K ((([Post Text], ()) -> IO ([Post Text], [Post Text]))
 -> K IO ([Post Text], ()) ([Post Text], [Post Text]))
-> (([Post Text], ()) -> IO ([Post Text], [Post Text]))
-> K IO ([Post Text], ()) ([Post Text], [Post Text])
forall a b. (a -> b) -> a -> b
$ \([Post Text]
s, ()) -> do
      os <- (AgentShard [Post Text] [Post Text] -> IO [Post Text])
-> [AgentShard [Post Text] [Post Text]] -> IO [[Post Text]]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse (AgentShard [Post Text] [Post Text] -> [Post Text] -> IO [Post Text]
`runSubShard` [Post Text]
s) [AgentShard [Post Text] [Post Text]]
shs
      pure ([], summary os)

-- | A fan-in summary honest by construction: the single output post is a
-- 'synthesis' of the supplied parent ids, so its ancestry cites every branch
-- output by exact reference.  The supplied function computes the body.
-- No branch outputs, or an empty body -> no posts (quiet).
--
-- The caller supplies one 'PostId' per branch output (in the order produced
-- by @concat oss@).  When ids are not available, pass @[]@ and the summary
-- falls back to a root post (no thread edge).
synthesisSummary ::
  Name ->
  [Name] ->
  [PostId] ->
  ([[Post Text]] -> Text) ->
  [[Post Text]] ->
  [Post Text]
synthesisSummary :: Text
-> [Text]
-> [PostId]
-> ([[Post Text]] -> Text)
-> [[Post Text]]
-> [Post Text]
synthesisSummary Text
who [Text]
audience [PostId]
parentIds [[Post Text]] -> Text
f [[Post Text]]
oss =
  case ([[Post Text]] -> [Post Text]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Post Text]]
oss, Text -> Text
T.strip ([[Post Text]] -> Text
f [[Post Text]]
oss)) of
    ([], Text
_) -> []
    ([Post Text]
_, Text
b) | Text -> Bool
T.null Text
b -> []
    ([Post Text]
ps, Text
b) ->
      let ids :: [PostId]
ids = Int -> [PostId] -> [PostId]
forall a. Int -> [a] -> [a]
take ([Post Text] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Post Text]
ps) [PostId]
parentIds
       in [if [PostId] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [PostId]
ids then Text -> [Text] -> Text -> Post Text
forall a. Text -> [Text] -> a -> Post a
mkPost Text
who [Text]
audience Text
b else Text -> [Text] -> [PostId] -> Text -> Post Text
forall b. Text -> [Text] -> [PostId] -> b -> Post b
synthesis Text
who [Text]
audience [PostId]
ids Text
b]