-- | Deterministic and probabilistic generators
module Numeric.Probability.Transition where

import qualified Numeric.Probability.Distribution as Dist

import qualified Control.Monad.Trans.Except as ME

import qualified Data.Map as Map

import qualified Data.List.HT as ListHT
import qualified Data.List as List
import Prelude hiding (map, maybe, id, )


-- * Transitions


-- | deterministic generator
type Change a = a -> a

-- | probabilistic generator
type T prob a = a -> Dist.T prob a


id :: (Num prob) => T prob a
id :: forall prob a. Num prob => T prob a
id = forall prob a. Num prob => T prob a
Dist.certainly


{- |
'map' maps a change function to the result of a transformation
('map' is somehow a lifted form of 'Dist.map')
The restricted type of @f@ results from the fact that the
argument to @t@ cannot be changed to @b@ in the result 'T' type.
-}
map :: (Num prob, Ord a) =>
   Change a -> T prob a -> T prob a
map :: forall prob a.
(Num prob, Ord a) =>
Change a -> T prob a -> T prob a
map Change a
f T prob a
t = forall prob b a.
(Num prob, Ord b) =>
(a -> b) -> T prob a -> T prob b
Dist.map Change a
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. T prob a
t


{- |
unfold a distribution of transitions into one transition

NOTE: The argument transitions must be independent
-}
unfold :: (Num prob, Ord a) =>
   Dist.T prob (T prob a) -> T prob a
unfold :: forall prob a. (Num prob, Ord a) => T prob (T prob a) -> T prob a
unfold T prob (T prob a)
d a
x = forall prob a. (Num prob, Ord a) => T prob (T prob a) -> T prob a
Dist.unfold (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a b. (a -> b) -> a -> b
$a
x) T prob (T prob a)
d)

{- |
Composition of transitions similar to 'Numeric.Probability.Monad.compose'
but with intermediate duplicate elimination.
-}
compose :: (Num prob, Ord a) =>
   [T prob a] -> T prob a
compose :: forall prob a. (Num prob, Ord a) => [T prob a] -> T prob a
compose = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\T prob a
acc T prob a
x a
v -> forall prob a. (Num prob, Ord a) => T prob a -> T prob a
Dist.norm (T prob a
acc a
v forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= T prob a
x)) forall (m :: * -> *) a. Monad m => a -> m a
return


untilLeft :: (Num prob, Ord a, Ord b) =>
   (a -> Dist.T prob (Either b a)) -> Dist.T prob a -> Dist.T prob b
untilLeft :: forall prob a b.
(Num prob, Ord a, Ord b) =>
(a -> T prob (Either b a)) -> T prob a -> T prob b
untilLeft a -> T prob (Either b a)
f =
   let go :: Map b prob -> T prob a -> T prob b
go Map b prob
final T prob a
dist =
          if forall (t :: * -> *) a. Foldable t => t a -> Bool
null (forall prob a. T prob a -> [(a, prob)]
Dist.decons T prob a
dist)
            then forall prob a. [(a, prob)] -> T prob a
Dist.Cons forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [(k, a)]
Map.toList Map b prob
final
            else
               case forall a b. [Either a b] -> ([a], [b])
ListHT.unzipEithers forall a b. (a -> b) -> a -> b
$
                    forall a b. (a -> b) -> [a] -> [b]
List.map (\(Either b a
e,prob
p) -> forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (\b
l -> forall a b. a -> Either a b
Left (b
l,prob
p)) (\a
r -> forall a b. b -> Either a b
Right (a
r,prob
p)) Either b a
e) forall a b. (a -> b) -> a -> b
$
                    forall prob a. T prob a -> [(a, prob)]
Dist.decons forall a b. (a -> b) -> a -> b
$ forall prob a. (Num prob, Ord a) => T prob a -> T prob a
Dist.norm forall a b. (a -> b) -> a -> b
$ T prob a
dist forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= a -> T prob (Either b a)
f of
                  ([(b, prob)]
newFinal, [(a, prob)]
stillActive) ->
                     Map b prob -> T prob a -> T prob b
go (forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
Map.unionWith forall a. Num a => a -> a -> a
(+) (forall k a. Ord k => (a -> a -> a) -> [(k, a)] -> Map k a
Map.fromListWith forall a. Num a => a -> a -> a
(+) [(b, prob)]
newFinal) Map b prob
final) forall a b. (a -> b) -> a -> b
$
                     forall prob a. [(a, prob)] -> T prob a
Dist.Cons [(a, prob)]
stillActive
   in  Map b prob -> T prob a -> T prob b
go forall k a. Map k a
Map.empty

{- |
In @fix $ \go a -> do ...; go xy@
any action after a 'go' is ignored.
-}
fix :: (Num prob, Ord a, Ord b) =>
   ((a -> ME.ExceptT a (Dist.T prob) b) ->
    (a -> ME.ExceptT a (Dist.T prob) b)) ->
   Dist.T prob a -> Dist.T prob b
fix :: forall prob a b.
(Num prob, Ord a, Ord b) =>
((a -> ExceptT a (T prob) b) -> a -> ExceptT a (T prob) b)
-> T prob a -> T prob b
fix (a -> ExceptT a (T prob) b) -> a -> ExceptT a (T prob) b
f =
   forall prob a b.
(Num prob, Ord a, Ord b) =>
(a -> T prob (Either b a)) -> T prob a -> T prob b
untilLeft forall a b. (a -> b) -> a -> b
$ \a
a ->
      case (a -> ExceptT a (T prob) b) -> a -> ExceptT a (T prob) b
f forall (m :: * -> *) e a. Monad m => e -> ExceptT e m a
ME.throwE a
a of
         ME.ExceptT T prob (Either a b)
m -> forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either forall a b. b -> Either a b
Right forall a b. a -> Either a b
Left) T prob (Either a b)
m


-- * Spreading changes into transitions

-- | functions to convert a list of changes into a transition
type SpreadC prob a = [Change a] -> T prob a

apply :: (Num prob) =>
   Change a -> T prob a
apply :: forall prob a. Num prob => Change a -> T prob a
apply Change a
f = forall prob a. Num prob => T prob a
id forall b c a. (b -> c) -> (a -> b) -> a -> c
. Change a
f


maybe :: (Num prob) => prob -> Change a -> T prob a
maybe :: forall prob a. Num prob => prob -> Change a -> T prob a
maybe prob
p Change a
f a
x = forall prob a. Num prob => prob -> a -> a -> T prob a
Dist.choose prob
p (Change a
f a
x) a
x

lift :: Dist.Spread prob a -> SpreadC prob a
lift :: forall prob a. Spread prob a -> SpreadC prob a
lift Spread prob a
s [Change a]
cs a
x = Spread prob a
s forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
List.map (forall a b. (a -> b) -> a -> b
$ a
x) [Change a]
cs

uniform :: (Fractional prob) => SpreadC prob a
uniform :: forall prob a. Fractional prob => SpreadC prob a
uniform  = forall prob a. Spread prob a -> SpreadC prob a
lift forall prob a. Fractional prob => Spread prob a
Dist.uniform

linear :: (Fractional prob) => SpreadC prob a
linear :: forall prob a. Fractional prob => SpreadC prob a
linear = forall prob a. Spread prob a -> SpreadC prob a
lift forall prob a. Fractional prob => Spread prob a
Dist.linear

normal :: (Floating prob) => SpreadC prob a
normal :: forall prob a. Floating prob => SpreadC prob a
normal   = forall prob a. Spread prob a -> SpreadC prob a
lift forall prob a. Floating prob => Spread prob a
Dist.normal

enum :: (RealFloat prob) => [Int] -> SpreadC prob a
enum :: forall prob a. RealFloat prob => [Int] -> SpreadC prob a
enum [Int]
xs  = forall prob a. Spread prob a -> SpreadC prob a
lift (forall prob a. Fractional prob => [Int] -> Spread prob a
Dist.enum [Int]
xs)

relative :: (RealFloat prob) => [prob] -> SpreadC prob a
relative :: forall prob a. RealFloat prob => [prob] -> SpreadC prob a
relative [prob]
xs  = forall prob a. Spread prob a -> SpreadC prob a
lift (forall prob a. Fractional prob => [prob] -> Spread prob a
Dist.relative [prob]
xs)


-- * Spreading transitions into transitions

-- | functions to convert a list of transitions into a transition
type SpreadT prob a = [T prob a] -> T prob a

liftT :: (Num prob, Ord a) =>
   Dist.Spread prob (T prob a) -> SpreadT prob a
liftT :: forall prob a.
(Num prob, Ord a) =>
Spread prob (T prob a) -> SpreadT prob a
liftT Spread prob (T prob a)
s = forall prob a. (Num prob, Ord a) => T prob (T prob a) -> T prob a
unfold forall b c a. (b -> c) -> (a -> b) -> a -> c
. Spread prob (T prob a)
s

uniformT :: (Fractional prob, Ord a) => SpreadT prob a
uniformT :: forall prob a. (Fractional prob, Ord a) => SpreadT prob a
uniformT  = forall prob a.
(Num prob, Ord a) =>
Spread prob (T prob a) -> SpreadT prob a
liftT forall prob a. Fractional prob => Spread prob a
Dist.uniform

linearT :: (Fractional prob, Ord a) => SpreadT prob a
linearT :: forall prob a. (Fractional prob, Ord a) => SpreadT prob a
linearT = forall prob a.
(Num prob, Ord a) =>
Spread prob (T prob a) -> SpreadT prob a
liftT forall prob a. Fractional prob => Spread prob a
Dist.linear

normalT :: (Floating prob, Ord a) => SpreadT prob a
normalT :: forall prob a. (Floating prob, Ord a) => SpreadT prob a
normalT   = forall prob a.
(Num prob, Ord a) =>
Spread prob (T prob a) -> SpreadT prob a
liftT forall prob a. Floating prob => Spread prob a
Dist.normal

enumT :: (RealFloat prob, Ord a) => [Int] -> SpreadT prob a
enumT :: forall prob a. (RealFloat prob, Ord a) => [Int] -> SpreadT prob a
enumT [Int]
xs  = forall prob a.
(Num prob, Ord a) =>
Spread prob (T prob a) -> SpreadT prob a
liftT (forall prob a. Fractional prob => [Int] -> Spread prob a
Dist.enum [Int]
xs)

relativeT :: (RealFloat prob, Ord a) => [prob] -> SpreadT prob a
relativeT :: forall prob a. (RealFloat prob, Ord a) => [prob] -> SpreadT prob a
relativeT [prob]
xs  = forall prob a.
(Num prob, Ord a) =>
Spread prob (T prob a) -> SpreadT prob a
liftT (forall prob a. Fractional prob => [prob] -> Spread prob a
Dist.relative [prob]
xs)