content: {-# LANGUAGE ScopedTypeVariables, BangPatterns #-}
  -- |Utility functions for working with 'Topic's. These functions are
  -- primarily combinators for fusing two 'Topic's in various ways.
  module Ros.Topic.Util where
  import Prelude hiding (dropWhile, filter, splitAt, mapM)
  import Control.Applicative
  import Control.Arrow ((***), second)
  import Control.Concurrent hiding (yield)
  import Control.Concurrent.STM
  import Control.Monad ((<=<), when, replicateM)
  import Control.Monad.IO.Class
  import Data.AdditiveGroup (AdditiveGroup, (^+^), (^-^), Sum(..))
  import Data.Monoid (Monoid)
  import Data.Sequence ((|>), viewl, ViewL(..))
  import qualified Data.Sequence as S
  import qualified Data.Foldable as F
  import Ros.Rate (rateLimiter)
  import Ros.Topic hiding (mapM_)
  
  -- |Produce an infinite list from a 'Topic'.
  toList :: Topic IO a -> IO [a]
  toList t0 = do c <- newChan
                 let feed t = do (x, t') <- runTopic t
                                 writeChan c x
                                 feed t'
                 _ <- forkIO $ feed t0
                 getChanContents c
  
  -- |Produce a 'Topic' from an infinite list.
  fromList :: Monad m => [a] -> Topic m a
  fromList (x:xs) = Topic $ return (x, fromList xs)
  fromList [] = error "Ran out of list elements"
  
  -- |Tee a 'Topic' into two duplicate 'Topic's. Each returned 'Topic'
  -- will receive all the values of the original 'Topic' while any
  -- side-effect produced by each step of the original 'Topic' will
  -- occur only once.
  -- 
  -- This version of @tee@ lazily pulls data from the original 'Topic'
  -- when it is first required by a consumer of either of the returned
  -- 'Topic's. This behavior is crucial when lazily consuming the data
  -- stream is preferred. For instance, using 'interruptible' with 'tee'
  -- will allow for a chunk of data to be abandoned before being fully
  -- consumed as long as neither consumer has forced its way too far
  -- down the stream.
  --
  -- This function is useful when two consumers must see all the same
  -- elements from a 'Topic'. If the 'Topic' was instead 'share'd, then
  -- one consumer might get the first value from the 'Topic' before the
  -- second consumer's buffer is created since buffer creation is lazy.
  tee :: Topic IO a -> IO (Topic IO a, Topic IO a)
  tee t0 = do c1 <- newTChanIO
              c2 <- newTChanIO
              signal <- newTVarIO True
              let feed c = do atomically $ do f <- isEmptyTChan c
                                              when f (writeTVar signal False)
                              atomically $ readTChan c
                  produce t = do atomically $ readTVar signal >>= flip when retry
                                 (x,t') <- runTopic t
                                 atomically $ writeTChan c1 x >>
                                              writeTChan c2 x >>
                                              writeTVar signal True
                                 produce t'
              _ <- forkIO $ produce t0
              return (repeatM (feed c1), repeatM (feed c2))
  
  -- |This version of @tee@ eagerly pulls data from the
  -- original 'Topic' as soon as it is available. This behavior is
  -- undesirable when lazily consuming the data stream is preferred. For
  -- instance, using 'interruptible' with 'teeEager' will likely not
  -- work well. However, 'teeEager' may have slightly better performance
  -- than 'tee'.
  teeEager :: Topic IO a -> IO (Topic IO a, Topic IO a)
  teeEager t = do c1 <- newChan
                  c2 <- newChan
                  let feed c = do x <- readChan c
                                  return (x, Topic $ feed c)
                  _ <- forkIO . forever . join $
                       (\x -> writeChan c1 x >> writeChan c2 x) <$> t
                  return (Topic $ feed c1, Topic $ feed c2)
  
  -- |Fan out one 'Topic' out to a number of duplicate 'Topic's, each of
  -- which will produce the same values. Side effects caused by the
  -- original 'Topic''s production will occur only once. This is useful
  -- when a known number of consumers must see exactly all the same
  -- elements.
  fan :: Int -> Topic IO a -> IO [Topic IO a]
  fan n t0 = do cs <- replicateM n newTChanIO
                signal <- newTVarIO True
                let feed c = do atomically $ do f <- isEmptyTChan c
                                                when f (writeTVar signal False)
                                atomically $ readTChan c
                    produce t = do atomically $ readTVar signal >>= flip when retry
                                   (x,t') <- runTopic t
                                   atomically $ mapM_ (flip writeTChan x) cs >>
                                                writeTVar signal True
                                   produce t'
                _ <- forkIO $ produce t0
                return $ map (repeatM . feed) cs
  
  -- |Make a 'Topic' shareable among multiple consumers. Each consumer
  -- of a Topic gets its own read buffer automatically as soon as it
  -- starts pulling items from the Topic. Without calling one of
  -- 'share', 'tee', or 'fan' on a Topic, the Topic's values will be
  -- split among all consumers (e.g. consumer /A/ gets half the values
  -- produced by the 'Topic', while consumer /B/ gets the other half
  -- with some unpredictable interleaving). Note that Topics returned by
  -- the @Ros.Node.subscribe@ are already shared.
  share :: Topic IO a -> IO (Topic IO a)
  share t0 = do cs <- newTVarIO [] -- A list for the individual client buffers
                signal <- newTVarIO True
                let addClient = atomically $ do cs0 <- readTVar cs
                                                c <- newTChan
                                                writeTVar cs (c:cs0)
                                                return c
                    feed c = do atomically $ do f <- isEmptyTChan c
                                                when f (writeTVar signal False)
                                atomically $ readTChan c
                    produce t = do atomically $ readTVar signal >>= flip when retry
                                   (x,t') <- runTopic t
                                   atomically $ do cs' <- readTVar cs
                                                   mapM_ (flip writeTChan x) cs'
                                                   writeTVar signal True
                                   produce t'
                _ <- forkIO $ produce t0
                return . Topic $ addClient >>= runTopic . repeatM . feed
  
  -- |The application @topicRate rate t@ runs 'Topic' @t@ no faster than
  -- @rate@ Hz.
  topicRate :: (Functor m, MonadIO m) => Double -> Topic m a -> Topic m a
  topicRate p t0 = Topic $ 
                   do delay <- liftIO $ rateLimiter p (return ())
                      (x,t') <- runTopic t0
                      let go t = Topic $ liftIO delay >> second go <$> runTopic t
                      return (x, go t')
  
  -- |Splits a 'Topic' into two 'Topic's: the elements of the first
  -- 'Topic' all satisfy the given predicate, while none of the elements
  -- of the second 'Topic' do.
  partition :: (a -> Bool) -> Topic IO a -> IO (Topic IO a, Topic IO a)
  partition p = fmap (filter p *** filter (not . p)) . tee
  
  -- |Returns a 'Topic' whose values are consecutive values from the
  -- original 'Topic'.
  consecutive :: Monad m => Topic m a -> Topic m (a,a)
  consecutive = metamorph startup
    where startup x = skip (go x)
          go x y    = yield (x,y) (go y)
  -- consecutive t = Topic $ do (x, t') <- runTopic t
  --                            runTopic $ go x t'
  --   where go x t' = Topic$ do (y, t'') <- runTopic t'
  --                             return ((x,y), go y t'')
  
  -- |Interleave two 'Topic's. Items from each component 'Topic' will be
  -- tagged with an 'Either' constructor and added to the combined
  -- 'Topic' as they become available.
  (<+>) :: Topic IO a -> Topic IO b -> Topic IO (Either a b)
  (<+>) t1 t2 = Topic $ do c <- newChan
                           let aux = do x <- readChan c
                                        return (x, Topic aux)
                               feed t = do (x,t') <- runTopic t
                                           writeChan c x
                                           feed t'
                           _ <- forkIO $ feed (fmap Left t1)
                           _ <- forkIO $ feed (fmap Right t2)
                           aux
  infixl 7 <+>
  
  -- |Returns a 'Topic' that produces a new pair every time either of
  -- the component 'Topic's produces a new value. The value of the
  -- other element of the pair will be the newest available value. The
  -- resulting 'Topic' will produce a new value at the rate of the
  -- faster component 'Topic', and may contain duplicate consecutive
  -- elements.
  everyNew :: Topic IO a -> Topic IO b -> Topic IO (a,b)
  everyNew t1 t2 = Topic $ warmup =<< runTopic (t1 <+> t2)
    where warmup (Left x, t)     = warmupR x =<< runTopic t
          warmup (Right y, t)    = warmupL y =<< runTopic t
          warmupR _ (Left x, t)  = warmupR x =<< runTopic t
          warmupR x (Right y, t) = return ((x,y), Topic $ runTopic t >>= go x y)
          warmupL _ (Right y, t) = warmupL y =<< runTopic t
          warmupL y (Left x, t)  = return ((x,y), Topic $ runTopic t >>= go x y)
          go _ y (Left x, t)     = return ((x,y), Topic $ runTopic t >>= go x y)
          go x _ (Right y, t)    = return ((x,y), Topic $ runTopic t >>= go x y)
  
  -- |Returns a 'Topic' that produces a new pair every time both of the
  -- component 'Topic's have produced a new value. The composite
  -- 'Topic' will produce pairs at the rate of the slower component
  -- 'Topic' consisting of the most recent value from each 'Topic'.
  bothNew :: Topic IO a -> Topic IO b -> Topic IO (a,b)
  bothNew t1 t2 = Topic $ warmup =<< runTopic (t1 <+> t2)
    where warmup (v,t) = go v =<< runTopic t
          go (Left _) (l@(Left _), t) = go l =<< runTopic t
          go (Left x) (Right y, t) = return ((x,y), Topic $ warmup =<< runTopic t)
          go (Right _) (r@(Right _), t) = go r =<< runTopic t
          go (Right y) (Left x, t) = return ((x,y), Topic $ warmup =<< runTopic t)
  
  -- |Returns a 'Topic' that produces a new pair every time a value of
  -- first topic produces a new value, followed by a new value from the
  -- second topic. This can be used for sampling the first topic with
  -- the second topic.
  firstThenSecond :: Topic IO a -> Topic IO b -> Topic IO (a,b)
  firstThenSecond t1 t2 = leftThenRight (t1 <+> t2)
  
  -- |Produces a value when a Left value is followed by a Right value.
  leftThenRight :: (Monad m) => Topic m (Either a b) -> Topic m (a,b)
  leftThenRight t1 = Topic $ warmup =<< runTopic t1
      where warmup (v,t) = go v =<< runTopic t
            go (Left x) (Right y, t) = return ((x,y), Topic $ warmup =<< runTopic t)
            go _ (x, t) = go x =<< runTopic t
  
  -- |Merge two 'Topic's into one. The items from each component
  -- 'Topic' will be added to the combined 'Topic' as they become
  -- available.
  merge :: Topic IO a -> Topic IO a -> Topic IO a
  merge t1 t2 = either id id <$> t1 <+> t2
  
  -- |Apply a function to each consecutive pair of elements from a 'Topic'.
  finiteDifference :: (Functor m, Monad m) => (a -> a -> b) -> Topic m a -> Topic m b
  finiteDifference f = fmap (uncurry f) . consecutive
  
  -- |Compute a running \"average\" of a 'Topic' using a user-provided
  -- normalization function applied to the sum of products. The
  -- arguments are a constat @alpha@ that is used to scale the current
  -- average, a constant @invAlpha@ used to scale the newest value, a
  -- function for adding two scaled values, a function for scaling
  -- input values, a function for normalizing the sum of scaled values,
  -- and finally the stream to average. Parameterizing over all the
  -- arithmetic to this extent allows for the use of denormalizing
  -- scaling factors, as might be used to keep all arithmetic
  -- integral. An example would be scaling the average by the integer
  -- 7, the new value by the integer 1, then normalizing by dividing
  -- the sum of scaled values by 8.
  weightedMeanNormalized :: Monad m =>
                            n -> n -> (b -> b -> c) -> (n -> a -> b) -> 
                            (c -> a) -> Topic m a -> Topic m a
  weightedMeanNormalized alpha invAlpha plus scale normalize = Topic . warmup
      where warmup = uncurry go <=< runTopic
            go avg t = do (x,t') <- runTopic t
                          let !avg' = normalize $ plus (scale alpha avg) 
                                                       (scale invAlpha x)
                          return (avg', Topic $ go avg' t')
  {-# INLINE weightedMeanNormalized #-}
  
  -- |Perform numerical integration of a 'Topic' using Simpson's rule
  -- applied at three consecutive points. This requires a function for
  -- adding values from the 'Topic', and a function for scaling values
  -- by a fractional number.
  simpsonsRule :: (Monad m, Fractional n) => 
                  (a -> a -> a) -> (n -> a -> a) -> Topic m a -> Topic m a
  simpsonsRule plus scale t0 = Topic $ do ([x,y], t') <- splitAt 2 t0
                                          go x y t'
    where go x y t = do (z,t') <- runTopic t
                        return (simpson x y z, Topic $ go y z t')
          simpson a mid b = scale c $ plus (plus a (scale 4 mid)) b
          c = 1 / 3
  {-# INLINE simpsonsRule #-}
  
  -- |Compute a running \"average\" of a 'Topic'. The application
  -- @weightedMean alpha plus scale t@ sums the product of @alpha@ and
  -- the current average with the product of @1 - alpha@ and the newest
  -- value produced by 'Topic' @t@. The addition and scaling operations
  -- are performed using the supplied @plus@ and @scale@ functions.
  weightedMean :: (Monad m, Num n) => 
                  n -> (a -> a -> a) -> (n -> a -> a) -> Topic m a -> Topic m a
  weightedMean alpha plus scale = weightedMean2 alpha (1 - alpha) plus scale
  {-# INLINE weightedMean #-}
  
  -- |Compute a running \"average\" of a 'Topic'. The application
  -- @weightedMean2 alpha invAlpha plus scale t@ sums the product of
  -- @alpha@ and the current average with the product of @invAlpha@ and
  -- the newest value produced by 'Topic' @t@. The addition and scaling
  -- operations are performed using the supplied @plus@ and @scale@
  -- functions.
  weightedMean2 :: Monad m =>
                   n -> n -> (a -> a -> a) -> (n -> a -> a) -> Topic m a -> Topic m a
  weightedMean2 alpha invAlpha plus scale = Topic . warmup
      where warmup = uncurry go <=< runTopic
            go avg t = do (x, t') <- runTopic t
                          let !savg = scale alpha avg
                              !sx = scale invAlpha x
                              !avg' = plus savg sx
                          return (avg', Topic $ go avg' t')
  {-# INLINE weightedMean2 #-}
  
  -- |Use a 'Topic' of functions to filter a 'Topic' of values. The
  -- application @filterBy t1 t2@ causes each function from 'Topic' @t1@
  -- to be applied to values produced by @t2@ until it returns
  -- 'True'. At that point, the 'filterBy' application produces the
  -- accepted value of the @t2@ and moves on to the next function from
  -- @t1@ which is applied to the rest of @t2@ in the same manner.
  filterBy :: Monad m => Topic m (a -> Bool) -> Topic m a -> Topic m a
  filterBy tf tx = Topic $ do (f, tf') <- runTopic tf
                              (x, tx') <- uncons $ dropWhile (not . f) tx
                              return (x, filterBy tf' tx')
  
  -- |Produce elements of the first 'Topic' no faster than elements of
  -- the second 'Topic' are produced.
  gate :: (Applicative m, Monad m) => Topic m a -> Topic m b -> Topic m a
  gate t1 t2 = const <$> t1 <*> t2
  
  -- |Flatten a 'Topic' of 'F.Foldable' values. For example, turn a
  -- @Topic m [a]@ of finite lists into a @Topic a@ by taking each
  -- element from each list in sequence.
  concats :: (Monad m, F.Foldable f) => Topic m (f a) -> Topic m a
  concats t = Topic $ do (x, t') <- runTopic t
                         F.foldr (\x' z -> return (x', Topic z)) 
                                 (runTopic $ concats t') 
                                 x
  
  -- |Flatten a 'Topic' of 'F.Foldable' values such that old values are
  -- discarded as soon as the original 'Topic' produces a new
  -- 'F.Foldable'.
  interruptible :: F.Foldable t => Topic IO (t a) -> Topic IO a
  interruptible s = Topic $
      do feeder <- newEmptyMVar         -- Active feeder thread
         latestItem <- newEmptyMVar     -- Next available item
         signal <- newEmptyMVar         -- Demand signal
         let feedItems ys = do ft <- tryTakeMVar feeder
                               maybe (return ()) killThread ft
                               t <- forkIO $ 
                                    F.traverse_ (\y -> takeMVar signal >> 
                                                       putMVar latestItem y) 
                                                ys
                               putMVar feeder t 
             watchForItems t = do (x,t') <- runTopic t
                                  feedItems x 
                                  watchForItems t'
             getAll = do putMVar signal ()
                         x <- takeMVar latestItem
                         return (x, Topic getAll)
         _ <- forkIO $ watchForItems s
         getAll
  
  -- |Pull elements from a 'Topic' in a new thread. This allows 'IO'
  -- 'Topic's to run at different rates even if they are consumed by a
  -- single thread.
  forkTopic :: Topic IO a -> IO (Topic IO a)
  forkTopic t = do c <- newChan
                   _ <- forkIO . forever . join $ fmap (writeChan c) t
                   let feed = Topic $ (\x -> (x,feed)) <$> readChan c
                   return feed
  
  -- |Sliding window over a 'Monoid'. @slidingWindow n t@ slides a
  -- window of width @n@ along 'Topic' @t@. As soon as at least @n@
  -- elements have been produced by @t@, the output 'Topic' starts
  -- producing the 'mconcat' of the elements in the window.
  slidingWindow :: (Monad m, Monoid a) => Int -> Topic m a -> Topic m a
  slidingWindow n = metamorph (fill S.empty)
    where fill w x
            | S.length w < n - 1 = skip . fill $ w |> x
            | otherwise = let w' = w |> x
                          in yield (F.fold w') (go w')
          go w x = let w' = dropOldest w |> x
                   in yield (F.fold w') (go w')
          dropOldest w = case viewl w of
                           EmptyL -> S.empty
                           _ :< w' -> w'
  
  -- |Sliding window over an 'AdditiveGroup'. @slidingWindowG n t@
  -- slides a window of width @n@ along 'Topic' @t@. As soon as at least
  -- @n@ elements have been produced by @t@, the output 'Topic' starts
  -- producing the total sum of the elements of the window. This
  -- function is more efficient than 'slidingWindow' because the group
  -- inverse operation is used to remove elements falling behind the
  -- window from the running sum.
  slidingWindowG :: (Monad m, AdditiveGroup a) => Int -> Topic m a -> Topic m a
  slidingWindowG n = metamorph (fill S.empty)
    where fill w x
            | S.length w < n - 1 = skip . fill $ w |> x
            | otherwise = let w' = w |> x
                              s = getSum . F.fold . fmap Sum $ w'
                          in yield s (go s w')
          go s w x = case viewl w of
                       EmptyL -> yield x $ go x (S.singleton x)
                       y :< w' -> let s' = s ^+^ x ^-^ y
                                  in yield s' $ go s' (w' |> x)
  
  -- |A way of pushing a monadic action into and along a 'Topic'. The
  -- application @topicOn proj inj trans t@ extracts a function from
  -- @trans@ that is then applied to the result of applying @proj@ to
  -- each value of 'Topic' @t@. The result of that application is
  -- supplied to the result of applying @inj@ to the same values from
  -- @t@ to produce a value for the output 'Topic'. A typical use case
  -- is projecting out a field from the original 'Topic' @t@ using
  -- @proj@ so that it may be modified by @trans@ and then injected back
  -- into the original structure using @inj@.
  topicOn :: (Applicative m, Monad m) =>
             (a -> b) -> (a -> c -> d) -> m (b -> m c) -> Topic m a -> Topic m d
  topicOn proj inj trans t = 
    Topic $ do f <- trans
               runTopic $ mapM (\x -> inj x `fmap` f (proj x)) t
  
  -- |@subsample n t@ subsamples topic 't' by dropping 'n' elements for
  -- every element produced by the result topic.
  subsample :: Monad m => Int -> Topic m b -> Topic m b
  subsample n = metamorph $ go n
    where go 0 x = yield x (go n)
          go i _ = skip (go (i - 1))
  