content: module Probability where
  
  import qualified Random
  import List (sort,sortBy,transpose)
  import Monad
  import Foreign (unsafePerformIO)
  
  import ListUtils
  import Show
  
  
  {- TO DO:
  
  * create export list
  
  * extend Dist by a constructor for continuous distributions:
  
    C (Float -> Float)
  
  * prove correctness of |||
  
  -}
  
  
  ------------------------------------------------------------------------------
  -- CONTENTS:
  --
  -- 0 AUXILIARY DEFINITIONS
  -- 1 DETERMINISTIC AND PROBABILISTIC VALUES
  -- 2 RANDOMIZED VALUES
  -- 3 DETERMINISTIC AND PROBABILISTIC GENERATORS
  -- 4 RANDOMIZED GENERATORS
  -- 5 ITERATORS AND SIMULATORS
  -- 6 TRACING
  ------------------------------------------------------------------------------
  
  
  ------------------------------------------------------------------------------
  -- 0 AUXILIARY DEFINITIONS
  --
  --   Event
  --   Probability
  ------------------------------------------------------------------------------
  
  -- 
  -- Events
  -- 
  type Event a = a -> Bool
  
  oneOf :: Eq a => [a] -> Event a
  oneOf = flip elem
  
  just :: Eq a => a -> Event a
  just = oneOf . singleton
  
  
  -- 
  -- Probabilities
  -- 
  newtype Probability = P ProbRep
  
  type ProbRep = Float
  
  precision :: Int
  precision = 1
  
  showPfix :: ProbRep -> String
  showPfix f | precision==0 = showR 3 (round (f*100))++"%"
             | otherwise    = showR (4+precision) (fromIntegral (round (f*100*d))/d)++"%"
               where d = 10^precision
  
  -- -- mixed precision
  -- -- 
  -- showP :: ProbRep -> String
  -- showP f | f>=0.1    = showR 3 (round (f*100))++"%"
  --         | otherwise = show (f*100)++"%"
   
  -- fixed precision
  -- 
  showP :: ProbRep -> String
  showP = showPfix
  
  
  instance Show Probability where
    show (P p) = showP p
  
  errorMargin :: ProbRep
  errorMargin = 0.00001
  
  
  --
  -- Monad composition
  --
  --  (>@>)  binary composition
  --  sequ   composition of a list of monadic functions
  --  
  (>@>) :: Monad m => (a -> m b) -> (b -> m c) -> a -> m c
  f >@> g = (>>= g) . f
  
  sequ :: Monad m => [a -> m a] -> a -> m a
  sequ = foldl (>@>) return
  
  
  
  ------------------------------------------------------------------------------
  -- 1 DETERMINISTIC AND PROBABILISTIC VALUES
  --
  --   Dist     probability disribution
  --   Spread   functions to convert a list of values into a distribution
  ------------------------------------------------------------------------------
  
  -- 
  -- Distributions
  -- 
  newtype Dist a = D {unD :: [(a,ProbRep)]}
  
  instance Monad Dist where
    return x = D [(x,1)]
    d >>= f  = D [(y,q*p) | (x,p) <- unD d, (y,q) <- unD (f x)]
    fail _   = D []
  
  -- note: mzero is a zero for >>= and a unit for mplus
  -- 
  instance MonadPlus Dist where
    mzero      = D []
    mplus d d' | isZero d || isZero d' = mzero
               | otherwise             = unfoldD $ choose 0.5 d d' 
  
  isZero :: Dist a -> Bool
  isZero (D d) = null d
  
  
  instance Functor Dist where
    fmap f (D d) = D [(f x,p) | (x,p) <- d]
  
  instance (Ord a,Eq a) => Eq (Dist a) where
    D xs == D ys = map fst (norm' xs)==map fst (norm' ys) && 
                     all (\((_,p),(_,q))->abs (p-q)<errorMargin) (zip xs ys)
  
  
  -- auxiliary functions for constructing and working with distributions
  -- 
  onD :: ([(a,ProbRep)] -> [(a,ProbRep)]) -> Dist a -> Dist a
  onD f  = D . f . unD
  
  sizeD :: Dist a -> Int
  sizeD = length . unD
  
  checkD :: Dist a -> Dist a
  checkD (D d) | abs (1-sumP d) < errorMargin = D d
               | otherwise = error ("Illegal distribution: total probability = "++show (sumP d))
  
  mkD :: [(a,ProbRep)] -> Dist a
  mkD = checkD . D
  
  sumP :: [(a,ProbRep)] -> ProbRep
  sumP = sum . map snd
  
  sortP :: [(a,ProbRep)] -> [(a,ProbRep)]
  sortP = sortBy (\x y->compare (snd y) (snd x))
  
  
  -- normalization = grouping
  -- 
  normBy ::  Ord a => (a -> a -> Bool) ->  Dist a -> Dist a
  normBy f = onD $ accumBy f . sort
  
  accumBy :: Num b => (a -> a -> Bool) -> [(a,b)] -> [(a,b)]
  accumBy f ((x,p):ys@((y,q):xs)) | f x y     = accumBy f ((x,p+q):xs)
                                  | otherwise = (x,p):accumBy f ys
  accumBy _ xs = xs
  
  norm ::  Ord a => Dist a -> Dist a
  norm = normBy (==)
  
  norm' :: Ord a => [(a,ProbRep)] -> [(a,ProbRep)]
  norm' = accumBy (==) . sort
  
  
  -- pretty printing
  -- 
  instance (Ord a,Show a) => Show (Dist a) where
    show (D []) = "Impossible"
    show (D xs) = concatMap (\(x,p)->showR w x++' ':showP p++"\n") (sortP (norm' xs))
                  where w = maximum (map (length.show.fst) xs)
  
  
  --
  -- Operations on distributions
  -- 
  
  -- product of independent distributions
  -- 
  joinWith :: (a -> b -> c) -> Dist a -> Dist b -> Dist c
  joinWith f (D d) (D d') = D [ (f x y,p*q) | (x,p) <- d, (y,q) <- d']
  
  prod :: Dist a -> Dist b -> Dist (a,b)
  prod = joinWith (,)
  
  
  -- distribution generators
  -- 
  type Spread a = [a] -> Dist a
  
  certainly :: Trans a
  certainly = return 
  
  impossible :: Dist a
  impossible = mzero
  
  choose :: ProbRep -> a -> a -> Dist a
  choose p x y = enum [p,1-p] [x,y]
  
  enum :: [ProbRep] -> Spread a
  enum ps xs = mkD $ zip xs ps
  
  enumPC :: [ProbRep] -> Spread a
  enumPC ps = enum (map (/100) ps)
  
  relative :: [Int] -> Spread a
  relative ns = enum (map (\n->fromIntegral n/fromIntegral (sum ns)) ns)
  
  shape :: (Float -> Float) -> Spread a
  shape _ [] = impossible
  shape f xs = scale (zip xs ps)
               where incr = 1 / fromIntegral ((length xs) - 1)
                     ps = map f (iterate (+incr) 0)
  
  linear :: Float -> Spread a
  linear c = shape (c*)
  
  uniform :: Spread a
  uniform = shape (const 1)
  
  negexp :: Spread a
  negexp = shape (\x -> exp (-x))
  
  normal :: Spread a
  normal = shape (normalCurve 0.5 0.5)
  
  normalCurve :: Float -> Float -> Float -> Float
  normalCurve mean stddev x = 1 / sqrt (2 * pi) * exp (-1/2 * u^2)
  	where u = (x - mean) / stddev
  
  
  -- extracting and mapping the domain of a distribution
  -- 
  extract :: Dist a -> [a]
  extract = map fst . unD
  
  mapD :: (a -> b) -> Dist a -> Dist b
  mapD = fmap
  
  
  -- unfold a distribution of distributions into one distribution
  -- 
  unfoldD :: Dist (Dist a) -> Dist a
  unfoldD (D d) = D [ (x,p*q) | (d',q) <- d, (x,p) <- unD d' ]
  
  
  -- conditional distribution
  -- 
  cond :: Dist Bool -> Dist a -> Dist a -> Dist a
  cond b d d' = unfoldD $ choose p d d'
                where P p = truth b
  
  truth :: Dist Bool -> Probability
  truth (D ((b,p):_)) = P (if b then p else 1-p)
  
  
  -- conditional probability
  -- 
  (|||) :: Dist a -> Event a -> Dist a
  (|||) = flip filterD
  
  
  -- filtering distributions
  --  
  data Select a = Case a | Other
                  deriving (Eq,Ord,Show)
  
  above :: Ord a => ProbRep -> Dist a -> Dist (Select a)
  above p (D d) = D (map (\(x,q)->(Case x,q)) d1++[(Other,sumP d2)])
                  where (d1,d2) = span (\(_,q)->q>=p) (sortP (norm' d))
  
  scale :: [(a,ProbRep)] -> Dist a
  scale xs = D (map (\(x,p)->(x,p/q)) xs)
             where q = sumP xs
                      
  filterD :: (a -> Bool) -> Dist a -> Dist a
  filterD p = scale . filter (p . fst) . unD
  
  
  -- selecting from distributions
  -- 
  selectP :: Dist a -> ProbRep -> a
  selectP (D d) p = scanP p d
  
  scanP :: ProbRep -> [(a,ProbRep)] -> a
  scanP p ((x,q):ps) | p<=q || null ps = x
                     | otherwise       = scanP (p-q) ps
  
  infix 8 ??
  
  (??) :: Event a -> Dist a -> Probability
  (??) p = P . sumP . filter (p . fst) . unD
  
   
  -- TO DO: generalize Float to arbitrary Num type
  -- 
  class ToFloat a where
    toFloat :: a -> Float
   
  instance ToFloat Float   where toFloat = id
  instance ToFloat Int     where toFloat = fromIntegral
  instance ToFloat Integer where toFloat = fromIntegral
  
  class FromFloat a where
    fromFloat :: Float -> a
  
  instance FromFloat Float   where fromFloat = id
  instance FromFloat Int     where fromFloat = round
  instance FromFloat Integer where fromFloat = round
    
  -- expected :: ToFloat a => Dist a -> Float
  -- expected = sum . map (\(x,p)->toFloat x*p) . unD
  
  class Expected a where
    expected :: a -> Float
  
  -- instance ToFloat a => Expected a where
  --   expected = toFloat
  instance Expected Float   where expected = id
  instance Expected Int     where expected = toFloat
  instance Expected Integer where expected = toFloat
  
  instance Expected a => Expected [a] where
    expected xs = sum (map expected xs) / toFloat (length xs)
  
  instance Expected a => Expected (Dist a) where
    expected = sum . map (\(x,p)->expected x*p) . unD
  
  instance Expected a => Expected (IO a) where
    expected r = expected (unsafePerformIO r)
  
  
  -- statistical analyses
  -- 
  variance :: Expected a => Dist a -> Float
  variance d@(D ps) = sum $ map (\(x,p)->p*sqr (expected x - ex)) ps
     where sqr x = x * x 
           ex    = expected d
  
  stddev :: Expected a => Dist a -> Float
  stddev = sqrt . variance
  
  
  
  ------------------------------------------------------------------------------
  -- 2 RANDOMIZED VALUES
  --
  --   R         random value
  --   RDist     random distribution
  ------------------------------------------------------------------------------
  
  
  --
  -- Random values
  --
  type R a = IO a
  
  printR :: Show a => R a -> R ()
  printR = (>>= print) 
  
  instance Show (IO a) where
    show _ = ""
  
  pick :: Dist a -> R a
  -- pick d = do {p <- Random.randomRIO (0,1); return (selectP p d)}
  pick d = Random.randomRIO (0,1) >>= return . selectP d
  
  
  --
  -- Randomized distributions
  --
  type RDist a = R (Dist a)
  
  rAbove :: Ord a => ProbRep -> RDist a -> RDist (Select a)
  rAbove p rd = do D d <- rd
                   let (d1,d2) = span (\(_,q)->q>=p) (sortP (norm' d))
                   return (D (map (\(x,q)->(Case x,q)) d1++[(Other,sumP d2)]))
  
  
  
  ------------------------------------------------------------------------------
  -- 3 DETERMINISTIC AND PROBABILISTIC GENERATORS
  --
  --   Change    deterministic generator
  --   Trans     probabilistic generator
  --   SpreadC   functions to convert a list of changes into a transition
  --   SpreadT   functions to convert a list of transitions into a transition
  ------------------------------------------------------------------------------
  
  -- 
  -- transitions
  -- 
  type Change a = a -> a
  
  type Trans a = a -> Dist a
  
  idT :: Trans a
  idT = certainlyT id
  
  
  -- mapT maps a change function to the result of a transformation
  -- (mapT is somehow a lifted form of mapD)
  -- The restricted type of f results from the fact that the
  -- argument to t cannot be changed to b in the result Trans type.
  -- 
  mapT :: Change a -> Trans a -> Trans a
  mapT f t = mapD f . t
  
  
  -- unfold a distribution of transitions into one transition
  -- 
  --   NOTE: The argument transitions must be independent
  -- 
  unfoldT :: Dist (Trans a) -> Trans a
  unfoldT (D d) x = D [ (y,p*q) | (f,p) <- d, (y,q) <- unD (f x) ]
  
  
  -- spreading changes into transitions
  -- 
  type SpreadC a = [Change a] -> Trans a
  
  certainlyT :: Change a -> Trans a
  certainlyT f = certainly . f
  -- certainlyT = (certainly .)
  -- certainlyT = maybeC 1
  
  maybeT :: ProbRep -> Change a -> Trans a
  maybeT p f = enumT [p,1-p] [f,id]
  
  liftC :: Spread a -> [Change a] -> Trans a
  liftC s cs x = s [f x | f <- cs]
  -- liftC s cs x = s $ map ($ x) cs
  
  uniformT  = liftC uniform
  normalT   = liftC normal
  linearT c = liftC (linear c)
  enumT xs  = liftC (enum xs)
  
  
  -- spreading transitions into transitions
  -- 
  type SpreadT a = [Trans a] -> Trans a
  
  liftT :: Spread (Trans a) -> [Trans a] -> Trans a
  liftT s = unfoldT . s
  
  uniformTT  = liftT uniform
  normalTT   = liftT normal
  linearTT c = liftT (linear c)
  enumTT xs  = liftT (enum xs)
  
   
  
  ------------------------------------------------------------------------------
  -- 4 RANDOMIZED GENERATORS
  --
  --   RChange   random change
  --   RTrans    random transition
  ------------------------------------------------------------------------------
  
  
  --
  -- Randomized changes
  --
  type RChange a = a -> R a
  
  random :: Trans a -> RChange a
  random t = pick . t
  -- random = (pick .)
  
  
  --
  -- Randomized transitions
  --
  type RTrans a = a -> RDist a
  type ApproxDist a = R [a]
  
  
  -- rDist converts a list of randomly generated values into
  --       a distribution by taking equal weights for all values
  --       
  rDist :: Ord a => [R a] -> RDist a
  rDist = fmap (norm . uniform) . sequence
  
  
  
  ------------------------------------------------------------------------------
  -- 5 ITERATION AND SIMULATION
  --
  -- Iterate   class defining *. 
  -- Sim       class defining ~.
  ------------------------------------------------------------------------------
  
  {-
  
  Naming convention:
  
  *   takes n :: Int and a generator and iterates the generator n times
  .   produces a single result
  ..  produces a trace
  ~   takes k :: Int [and n :: Int] and a generator and simulates 
      the [n-fold repetition of the] generator k times
  
  n *.  t   iterates t and produces a distribution
  n *.. t   iterates t and produces a trace
  
  k     ~.  t   simulates t and produces a distribution
  (k,n) ~*. t   simulates the n-fold repetition of t and produces a distribution
  (k,n) ~.. t   simulates the n-fold repetition of t and produces a trace
  
  -}
  
  -- Iteration captures three iteration strategies:
  -- iter builds an n-fold composition of a (randomized) transition 
  -- while and until implement conditional repetitions
  --
  -- The class Iterate allows the overloading of iteration for different
  -- kinds of generators, namely transitions and random changes:
  -- 
  --   Trans   a = a -> Dist a    ==>   c = Dist
  --   RChange a = a -> R a       ==>   c = R = IO
  -- 
  class Iterate c where
    (*.)  :: Int -> (a -> c a) -> (a -> c a)
    while :: (a -> Bool) -> (a -> c a) -> (a -> c a)
    until :: (a -> Bool) -> (a -> c a) -> (a -> c a)
    until p = while (not.p)
  
  infix 8 *.
  
  -- iteration of transitions
  -- 
  instance Iterate Dist where
    n *. t = head . (n *.. t)
    while p t x = if p x then t x >>= while p t else certainly x
  
  -- iteration of random changes
  -- 
  instance Iterate IO where
    n *. r = (>>= return . head) . rWalk n r
    while p t x = do {l <- t x; if p l then while p t l else return l}
  
  
  
  -- Simulation means to repeat a random chage many times and
  -- to accumulate all results into a distribution. Therefore,
  -- simulation can be regarded as an approximation of distributions
  -- through randomization.
  --
  -- The Sim class contains two functions: 
  -- 
  --   ~.   returns the final randomized transition 
  --   ~..  returns the whole trace
  -- 
  -- The Sim class allows the overloading of simulation for different
  -- kinds of generators, namely transitions and random changes:
  -- 
  --   Trans   a = a -> Dist a   ==>   c = Dist
  --   RChange a = a -> R a      ==>   c = R = IO
  -- 
  class Sim c where 
    (~.)  :: Ord a => Int       -> (a -> c a) -> RTrans a
    (~..) :: Ord a => (Int,Int) -> (a -> c a) -> RExpand a
    (~*.) :: Ord a => (Int,Int) -> (a -> c a) -> RTrans a
  
  infix 6 ~.
  infix 6 ~..
  
  -- simulation for transitions
  -- 
  instance Sim Dist where
    (~.)  x = (~.)  x . random
    (~..) x = (~..) x . random
    (~*.) x = (~*.) x . random
   
  
  -- simulation for random changes
  -- 
  instance Sim IO where  
    (~.)     n  t = rDist . replicate n . t
    (~..) (k,n) t = mergeTraces . replicate k . rWalk n t
    (~*.) (k,n) t = k ~. n *. t
  
  infix 8 ~*.
  
  --(~*.) :: (Iterate c,Sim c,Ord a) => (Int,Int) -> (a -> c a) -> RTrans a
  --(k,n) ~*. t = 
  
  
  ------------------------------------------------------------------------------
  -- 7 TRACING
  --
  --   (R)Trace
  --   (R)Space
  --   (R)Walk
  --   (R)Expand
  ------------------------------------------------------------------------------
  
  type Trace a  = [a]
  type Space a  = Trace (Dist a)
  type Walk a   = a -> Trace a
  type Expand a = a -> Space a
  
  
  -- >>: composes the result of a transition with a space 
  -- (transition is composed on the left)
  -- 
  -- (a -> m a) -> (a -> [m a]) -> (a -> [m a])
  (>>:) :: Trans a -> Expand a -> Expand a
  f >>: g = \x -> let ds@(D d:_)=g x in
                      D [ (z,p*q) | (y,p) <- d, (z,q) <- unD (f y)]:ds
  
  infix 6 >>:
  
  -- walk is a bounded version of the predefined function iterate
  -- 
  walk :: Int -> Change a -> Walk a
  walk n f = take n . iterate f
  
  -- *.. is identical to *., but returns the list of all intermediate
  --     distributions
  --     
  (*..) :: Int -> Trans a -> Expand a
  0 *.. _ = singleton . certainly
  1 *.. t = singleton . t
  n *.. t = t >>: (n-1) *.. t
  
  infix 8 *..
  
  
  type RTrace a  = R (Trace a)
  type RSpace a  = R (Space a)
  type RWalk a   = a -> RTrace a
  type RExpand a = a -> RSpace a
  
  --          (a -> m a) -> (a -> m [a]) -> (a -> m [a])
  composelR :: RChange a -> RWalk a -> RWalk a
  composelR f g x = do {rs@(r:_) <- g x; s <- f r; return (s:rs)}
  
  
  -- rWalk computes a list of values by
  --       randomly selecting one value from a distribution in each step.
  -- 
  rWalk :: Int -> RChange a -> RWalk a
  rWalk 0 _ = return . singleton
  rWalk 1 t = (>>= return . singleton) . t
  rWalk n t = composelR t (rWalk (n-1) t)
  
  
  -- mergeTraces converts a list of RTraces, into a list of randomized 
  --             distributions, i.e., an RSpace, by creating a randomized
  --             distribution for each list position across all traces
  --
  mergeTraces :: Ord a => [RTrace a] -> RSpace a
  mergeTraces = fmap (zipListWith (norm . uniform)) . sequence
                where
                  zipListWith :: ([a] -> b) -> [[a]] -> [b]
                  zipListWith f = map f . transpose
  
  {-
  
  LAWS
  
    const . pick = random . const
  
  -}
  
  