-- |
-- 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
      WContext(..),
      WContextLimit(..),

      -- * Distribute
      WDistribute(..)
    )
    where

import qualified Data.Map as M
import Control.Arrow
import Control.Wire.Classes
import Control.Wire.TimedMap
import Control.Wire.Types
import Data.AdditiveGroup
import Data.Either


-- | 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.

class Arrow (>~) => WContext (>~) where
    context :: Ord k => Wire e (>~) (a, k) b -> Wire e (>~) (a, k) b

instance Monad m => WContext (Kleisli m) where
    context w0 = context' M.empty
        where
        --context' :: Ord k => Map k (Wire e (Kleisli m) (a, k) b) -> Wire e (Kleisli m) (a, k) b
        context' ctxs' =
            WmGen $ \(x', ctx) -> do
                let w' = M.findWithDefault w0 ctx ctxs'
                (mx, w) <- toGenM w' (x', ctx)
                let ctxs = M.insert ctx w ctxs'
                return (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.

class Arrow (>~) => WContextLimit t (>~) | (>~) -> t where
    contextLimit :: Ord k => Wire e (>~) (a, k) b -> Wire e (>~) ((a, k), t) b

instance (AdditiveGroup t, MonadClock t m, Ord t) => WContextLimit t (Kleisli m) where
    contextLimit w0 = context' tmEmpty
        where
        --context' :: Ord k => TimedMap t k (Wire e (Kleisli m) (a, k) b) -> Wire e (Kleisli m) ((a, k), t) b
        context' ctxs' =
            WmGen $ \((x', ctx), maxAge) -> do
                t <- getTime
                let w' = tmFindWithDefault w0 ctx ctxs'
                (mx, w) <- toGenM w' (x', ctx)

                let ctxs = tmLimitAge (t ^-^ maxAge) . tmInsert t ctx w $ ctxs'
                return (mx, context' ctxs)


-- | Distribute the input signal over the given wires, evolving each of
-- them individually.  Discards all inhibited signals.
--
-- * Depends: as strict as the strictest subwire.

class Arrow (>~) => WDistribute (>~) where
    distribute :: [Wire e (>~) a b] -> Wire e (>~) a [b]

instance Monad m => WDistribute (Kleisli m) where
    distribute ws' =
        WmGen $ \x' -> do
            res <- mapM (\w' -> toGenM w' x') ws'
            let (mxs, ws) = first rights . unzip $ res
            return (Right mxs, distribute ws)