-- |
-- Module:     Control.Wire.Trans.Combine
-- Copyright:  (c) 2011 Ertugrul Soeylemez
-- License:    BSD3
-- Maintainer: Ertugrul Soeylemez <es@ertes.de>
--
-- Wire transformers for combining wires.

module Control.Wire.Trans.Combine
    ( -- * Context-sensitive evolution
      context,
      contextLimit,

      -- * Distribute
      distribute
    )
    where

import qualified Data.Map as M
import qualified Data.Set as S
import Control.Arrow
import Control.Wire.Classes
import Control.Wire.Tools
import Control.Wire.Types
import Data.Either
import Data.Map (Map)
import Data.Set (Set)


-- | Make the given wire context-sensitive.  The right signal is a
-- context and the wire will evolve individually for each context.
--
-- * Depends: Like context wire (left), current instant (right).
-- * Inhibits: Like context wire.

context ::
    forall a b e k (>~).
    (ArrowApply (>~), ArrowChoice (>~), Ord k)
    => Wire e (>~) a b -> Wire e (>~) (a, k) b
context w0 = context' M.empty
    where
    context' :: Map k (Wire e (>~) a b) -> Wire e (>~) (a, k) b
    context' ctxs' =
        mkGen $ proc (x', ctx) -> do
            let w' = M.findWithDefault w0 ctx ctxs'
            (mx, w) <- toGen w' -<< x'
            let ctxs = M.insert ctx w ctxs'
            returnA -< (mx, context' ctxs)


-- | Same as 'context', but with a time limit.  The third signal
-- specifies a maximum age.  Contexts not used for longer than the
-- maximum age are forgotten.
--
-- * Depends: Like context wire (left), current instant (right).
-- * Inhibits: Like context wire.

contextLimit ::
    forall a b e k t (>~).
    (ArrowApply (>~), ArrowClock (>~), Num t, Ord k, Ord t, Time (>~) ~ t)
    => Wire e (>~) a b -> Wire e (>~) ((a, k), t) b
contextLimit w0 = context' M.empty M.empty
    where
    context' ::
        Map k (Wire e (>~) a b, t)
        -> Map t (Set k)
        -> Wire e (>~) ((a, k), t) b
    context' ctxs'' hist'' =
        mkGen $ proc ((x', ctx), maxAge) -> do
            t <- arrTime -< ()
            let (w', t') = M.findWithDefault (w0, t) ctx ctxs''
            (mx, w) <- toGen w' -<< x'

            let ctxs' = M.insert ctx (w, t) ctxs''
                hist' =
                    M.insertWith' S.union t (S.singleton ctx) .
                    M.update (\s' -> let s = S.delete ctx s'
                                     in if S.null s then Nothing else Just s) t' $
                    hist''

                (ctxs, hist) =
                    let (delMap, hist) = M.split (t - maxAge) hist'
                        dels = M.fromDistinctAscList . map (, ()) .
                               S.toAscList . S.unions . M.elems $ delMap
                    in (ctxs' M.\\ dels, hist)
            returnA -< (mx, context' ctxs hist)


-- | Distribute the input signal over the given wires, evolving each of
-- them individually.  Collects produced outputs.
--
-- Note: This wire transformer discards all inhibited signals.
--
-- * Depends: as strict as the strictest subwire.

distribute ::
    ArrowApply (>~) =>
    [Wire e (>~) a b] -> Wire e (>~) a [b]
distribute ws' =
    mkGen $ proc x' -> do
        (mxs, ws) <-
            first rights . unzip ^<<
            distA (map toGen ws') -<< x'
        returnA -< (Right mxs, distribute ws)