module Dice where

import Data.List as L
import Data.Ord (comparing)
import Data.Ratio

newtype Prob a = Prob { getProb :: [(a,Rational)] } deriving Show

instance Functor Prob where
    fmap f (Prob xs) = Prob $ map (\(x,p) -> (f x, p)) xs

flatten :: Prob (Prob a) -> Prob a
flatten (Prob xs) = Prob $ concat $ map multAll xs
    where multAll (Prob xxs, p) = map (\(x,r) -> (x,p*r)) xxs

instance Monad Prob where
    return x = Prob [(x,1%1)]
    m >>= f = flatten (fmap f m)
    fail _ = Prob []

dice :: Prob Int
dice = Prob [(1, 1%6), (2, 1%6), (3, 1%6), (4, 1%6), (5, 1%6), (6, 1%6)]

roll6 :: Prob Bool
roll6 = do
        a <- dice
        b <- dice
        c <- dice
        d <- dice
        e <- dice
        f <- dice
        return (all (== 1) [a,b,c,d,e,f])

distinct :: (Ord a) => Prob a -> Prob a
distinct xs = let x = L.groupBy (\x y -> fst x == fst y) . L.sortBy (comparing fst) . getProb $ xs
                     in Prob $ map (L.foldl1' addTup) x
                         where addTup (a, p) (_, r) = (a, p+r)