-- | Probability-related types.
--
-- TODO instances for serialization and further stuff
-- TODO vector instances

module Statistics.Probability where

import Control.DeepSeq
import Control.Lens
import Data.Aeson
import Data.Char (chr)
import Data.Vector.Unboxed.Deriving
import Data.Vector.Unboxed (Unbox)
import GHC.Generics(Generic)
import Numeric.Log

import Algebra.Structure.Semiring
import Numeric.LogDomain
import Numeric.Limits



data IsNormalized = Normalized | NotNormalized



-- * Probability in linear space

-- | @Prob@ wraps a @Double@ that encodes probabilities. If @Prob@ is tagged as
-- @Normalized@, the contained values are in the range @[0,...,1]@, otherwise
-- they are in the range @[0,...,∞]@.

newtype Probability (n ∷ IsNormalized) x = Prob { Probability n x -> x
getProb ∷ x }
  deriving (Probability n x -> Probability n x -> Bool
(Probability n x -> Probability n x -> Bool)
-> (Probability n x -> Probability n x -> Bool)
-> Eq (Probability n x)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall (n :: IsNormalized) x.
Eq x =>
Probability n x -> Probability n x -> Bool
/= :: Probability n x -> Probability n x -> Bool
$c/= :: forall (n :: IsNormalized) x.
Eq x =>
Probability n x -> Probability n x -> Bool
== :: Probability n x -> Probability n x -> Bool
$c== :: forall (n :: IsNormalized) x.
Eq x =>
Probability n x -> Probability n x -> Bool
Eq,Eq (Probability n x)
Eq (Probability n x)
-> (Probability n x -> Probability n x -> Ordering)
-> (Probability n x -> Probability n x -> Bool)
-> (Probability n x -> Probability n x -> Bool)
-> (Probability n x -> Probability n x -> Bool)
-> (Probability n x -> Probability n x -> Bool)
-> (Probability n x -> Probability n x -> Probability n x)
-> (Probability n x -> Probability n x -> Probability n x)
-> Ord (Probability n x)
Probability n x -> Probability n x -> Bool
Probability n x -> Probability n x -> Ordering
Probability n x -> Probability n x -> Probability n x
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall (n :: IsNormalized) x. Ord x => Eq (Probability n x)
forall (n :: IsNormalized) x.
Ord x =>
Probability n x -> Probability n x -> Bool
forall (n :: IsNormalized) x.
Ord x =>
Probability n x -> Probability n x -> Ordering
forall (n :: IsNormalized) x.
Ord x =>
Probability n x -> Probability n x -> Probability n x
min :: Probability n x -> Probability n x -> Probability n x
$cmin :: forall (n :: IsNormalized) x.
Ord x =>
Probability n x -> Probability n x -> Probability n x
max :: Probability n x -> Probability n x -> Probability n x
$cmax :: forall (n :: IsNormalized) x.
Ord x =>
Probability n x -> Probability n x -> Probability n x
>= :: Probability n x -> Probability n x -> Bool
$c>= :: forall (n :: IsNormalized) x.
Ord x =>
Probability n x -> Probability n x -> Bool
> :: Probability n x -> Probability n x -> Bool
$c> :: forall (n :: IsNormalized) x.
Ord x =>
Probability n x -> Probability n x -> Bool
<= :: Probability n x -> Probability n x -> Bool
$c<= :: forall (n :: IsNormalized) x.
Ord x =>
Probability n x -> Probability n x -> Bool
< :: Probability n x -> Probability n x -> Bool
$c< :: forall (n :: IsNormalized) x.
Ord x =>
Probability n x -> Probability n x -> Bool
compare :: Probability n x -> Probability n x -> Ordering
$ccompare :: forall (n :: IsNormalized) x.
Ord x =>
Probability n x -> Probability n x -> Ordering
$cp1Ord :: forall (n :: IsNormalized) x. Ord x => Eq (Probability n x)
Ord,Int -> Probability n x -> ShowS
[Probability n x] -> ShowS
Probability n x -> String
(Int -> Probability n x -> ShowS)
-> (Probability n x -> String)
-> ([Probability n x] -> ShowS)
-> Show (Probability n x)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall (n :: IsNormalized) x.
Show x =>
Int -> Probability n x -> ShowS
forall (n :: IsNormalized) x. Show x => [Probability n x] -> ShowS
forall (n :: IsNormalized) x. Show x => Probability n x -> String
showList :: [Probability n x] -> ShowS
$cshowList :: forall (n :: IsNormalized) x. Show x => [Probability n x] -> ShowS
show :: Probability n x -> String
$cshow :: forall (n :: IsNormalized) x. Show x => Probability n x -> String
showsPrec :: Int -> Probability n x -> ShowS
$cshowsPrec :: forall (n :: IsNormalized) x.
Show x =>
Int -> Probability n x -> ShowS
Show,ReadPrec [Probability n x]
ReadPrec (Probability n x)
Int -> ReadS (Probability n x)
ReadS [Probability n x]
(Int -> ReadS (Probability n x))
-> ReadS [Probability n x]
-> ReadPrec (Probability n x)
-> ReadPrec [Probability n x]
-> Read (Probability n x)
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
forall (n :: IsNormalized) x. Read x => ReadPrec [Probability n x]
forall (n :: IsNormalized) x. Read x => ReadPrec (Probability n x)
forall (n :: IsNormalized) x.
Read x =>
Int -> ReadS (Probability n x)
forall (n :: IsNormalized) x. Read x => ReadS [Probability n x]
readListPrec :: ReadPrec [Probability n x]
$creadListPrec :: forall (n :: IsNormalized) x. Read x => ReadPrec [Probability n x]
readPrec :: ReadPrec (Probability n x)
$creadPrec :: forall (n :: IsNormalized) x. Read x => ReadPrec (Probability n x)
readList :: ReadS [Probability n x]
$creadList :: forall (n :: IsNormalized) x. Read x => ReadS [Probability n x]
readsPrec :: Int -> ReadS (Probability n x)
$creadsPrec :: forall (n :: IsNormalized) x.
Read x =>
Int -> ReadS (Probability n x)
Read,(forall x. Probability n x -> Rep (Probability n x) x)
-> (forall x. Rep (Probability n x) x -> Probability n x)
-> Generic (Probability n x)
forall x. Rep (Probability n x) x -> Probability n x
forall x. Probability n x -> Rep (Probability n x) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall (n :: IsNormalized) x x.
Rep (Probability n x) x -> Probability n x
forall (n :: IsNormalized) x x.
Probability n x -> Rep (Probability n x) x
$cto :: forall (n :: IsNormalized) x x.
Rep (Probability n x) x -> Probability n x
$cfrom :: forall (n :: IsNormalized) x x.
Probability n x -> Rep (Probability n x) x
Generic)

instance (NFData x) ⇒ NFData (Probability n x)

derivingUnbox "Probability"
  [t| forall n x. Unbox x ⇒ Probability n x → x |]  [| getProb |]  [| Prob |]

deriving instance (Enum       x) ⇒ Enum       (Probability n x)
deriving instance (Num        x) ⇒ Num        (Probability n x)
deriving instance (Fractional x) ⇒ Fractional (Probability n x)
deriving instance (Floating   x) ⇒ Floating   (Probability n x)
deriving instance (Real       x) ⇒ Real       (Probability n x)
deriving instance (RealFrac   x) ⇒ RealFrac   (Probability n x)
deriving instance (RealFloat  x) ⇒ RealFloat  (Probability n x)

instance ToJSON x ⇒ ToJSON (Probability n x) where
  toJSON :: Probability n x -> Value
toJSON = x -> Value
forall a. ToJSON a => a -> Value
toJSON (x -> Value) -> (Probability n x -> x) -> Probability n x -> Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Probability n x -> x
forall (n :: IsNormalized) x. Probability n x -> x
getProb

instance FromJSON x ⇒ FromJSON (Probability n x) where
  parseJSON :: Value -> Parser (Probability n x)
parseJSON = (x -> Probability n x) -> Parser x -> Parser (Probability n x)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap x -> Probability n x
forall (n :: IsNormalized) x. x -> Probability n x
Prob (Parser x -> Parser (Probability n x))
-> (Value -> Parser x) -> Value -> Parser (Probability n x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Value -> Parser x
forall a. FromJSON a => Value -> Parser a
parseJSON

instance (Num r) ⇒ Semiring (Probability n r) where
  plus :: Probability n r -> Probability n r -> Probability n r
plus  = Probability n r -> Probability n r -> Probability n r
forall a. Num a => a -> a -> a
(+)
  times :: Probability n r -> Probability n r -> Probability n r
times = Probability n r -> Probability n r -> Probability n r
forall a. Num a => a -> a -> a
(*)
  zero :: Probability n r
zero = Probability n r
0
  one :: Probability n r
one  = Probability n r
1
  {-# Inline plus  #-}
  {-# Inline times #-}
  {-# Inline zero  #-}
  {-# Inline one   #-}

-- | Turns a value into a normalized probability. @error@ if the value is not
-- in the range @[0,...,1]@.

prob ∷ (Ord x, Num x, Show x) ⇒ x → Probability Normalized x
prob :: x -> Probability 'Normalized x
prob x
x
  | x
x x -> x -> Bool
forall a. Ord a => a -> a -> Bool
>= x
0 Bool -> Bool -> Bool
&& x
x x -> x -> Bool
forall a. Ord a => a -> a -> Bool
<= x
1 = x -> Probability 'Normalized x
forall (n :: IsNormalized) x. x -> Probability n x
Prob x
x
  | Bool
otherwise        = String -> Probability 'Normalized x
forall a. HasCallStack => String -> a
error (String -> Probability 'Normalized x)
-> String -> Probability 'Normalized x
forall a b. (a -> b) -> a -> b
$ x -> String
forall a. Show a => a -> String
show x
x String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" not in range of [0,...,1]"
{-# Inline prob #-}

-- | Simple wrapper around @Probability@ that fixes non-normalization.

prob' ∷ (Ord x, Num x, Show x) ⇒ x → Probability NotNormalized x
prob' :: x -> Probability 'NotNormalized x
prob' = x -> Probability 'NotNormalized x
forall (n :: IsNormalized) x. x -> Probability n x
Prob
{-# Inline prob' #-}

-- | This simple function represents probabilities with characters between '0'
-- @0.0 -- 0.05@ up to '9' @0.85 -- 0.95@ and finally '*' for @>0.95@.

probabilityToChar ∷ (Num k, RealFrac k) ⇒ Probability Normalized k → Char
probabilityToChar :: Probability 'Normalized k -> Char
probabilityToChar (Prob k
p')
  | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
10 = Char
'*'
  | Bool
otherwise = Int -> Char
chr (Int -> Char) -> Int -> Char
forall a b. (a -> b) -> a -> b
$ Int
48 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i
  where p :: k
p = k -> k -> k
forall a. Ord a => a -> a -> a
max k
0.0 (k -> k) -> k -> k
forall a b. (a -> b) -> a -> b
$ k -> k -> k
forall a. Ord a => a -> a -> a
min k
p' k
1.0
        i :: Int
i = k -> Int
forall a b. (RealFrac a, Integral b) => a -> b
round (k -> Int) -> k -> Int
forall a b. (a -> b) -> a -> b
$ k
p k -> k -> k
forall a. Num a => a -> a -> a
* k
10


-- -- * Probability in log space. A number of operations internally cast to @Log@
-- -- from @log-domain@, but the values themselves are *not* newtype-wrapped @Log
-- -- x@ values. This is because we want to be *explicit* that these are
-- -- log-probabilities.
-- --
-- -- @Log@ numbers in cases like @fromIntegral 1 :: Log Double@ are treated as
-- -- not being in the log-domain, hence @fromIntegral performs a @log@
-- -- operations.
-- 
-- newtype LogProb (n ∷ IsNormalized) x = LogProb { getLogProb ∷ x }
--   deriving (Eq,Ord,Show)
-- 
-- derivingUnbox "LogProb"
--   [t| forall n x. Unbox x ⇒ LogProb n x → x |]  [| getLogProb |]  [| LogProb |]
-- 
-- instance (Precise x, RealFloat x) ⇒ Num (LogProb n x) where
--   (+) = withLog2 (+)
--   {-# Inline (+) #-}
--   (*) = withLog2 (*)
--   {-# Inline (*) #-}
--   abs = withLog1 abs
--   {-# Inline abs #-}
--   signum = withLog1 signum
--   {-# Inline signum #-}
--   fromInteger = LogProb . fromInteger
--   {-# Inline fromInteger #-}
--   negate = withLog1 negate
--   {-# Inline negate #-}
--   (-) = withLog2 (-)
--   {-# Inline (-) #-}
-- 
-- instance (Precise r, RealFloat r, Num r) ⇒ SemiRing (LogProb n r) where
--   srplus = (+)
--   {-# Inline srplus #-}
--   srmul  = (*)
--   {-# Inline srmul #-}
--   srzero = 0
--   {-# Inline srzero #-}
--   srone  = 1
--   {-# Inline srone #-}
-- 
-- instance (Num d, Fractional d) ⇒ NumericLimits (LogProb n d) where
--   minFinite = LogProb 0
--   maxFinite = LogProb (1/0)
-- 
-- withLog1 ∷ (Log x → Log y) → LogProb n x → LogProb n y
-- withLog1 op (LogProb x) = LogProb . ln $ op (Exp x)
-- {-# Inline withLog1 #-}
-- 
-- withLog2 ∷ (Log x → Log y → Log z) → LogProb n x → LogProb n y → LogProb n z
-- withLog2 op (LogProb x) (LogProb y) = LogProb . ln $ op (Exp x) (Exp y)
-- {-# Inline withLog2 #-}
-- 
-- 
-- -- * Conversion between probability in linear and log-space.
-- 
-- -- | Turn probability into log-probability.
-- 
-- p2lp ∷ (Floating x) ⇒ Prob n x → LogProb n x
-- p2lp (Prob x) = LogProb $ log x
-- {-# Inline p2lp #-}
-- 
-- -- | Turn log-probability into probability.
-- 
-- lp2p ∷ (Floating x) ⇒ LogProb n x → Prob n x
-- lp2p (LogProb x) = Prob $ exp x
-- {-# Inline lp2p #-}
-- 
-- -- | An isomorphism between @Prob@ and @LogProb@.
-- 
-- aslp ∷ (Floating x) ⇒ Iso' (Prob n x) (LogProb n x)
-- aslp = iso p2lp lp2p
-- {-# Inline aslp #-}