module Rolling.Prob where

import GHC.Exts (groupWith)

-- a probability expressed as a ratio (i.e. @P 1 2@ is a 50/50 chance)
data P = P Int Int deriving (Eq, Ord)

instance Show P where
  show (P n t) = show n ++ " in " ++ show t

times :: P -> P -> P
P a at `times` P b bt = P (a * b) (at * bt)

combine :: [P] -> P
combine ps =
  let tots = [ t | P _ t <- ps ]
      newTot = foldr lcm 1 tots
      newNums = [ if t == newTot then n else error "FIXME" | P n t <- ps ]
  in P (sum newNums) newTot


-- a Prob monad that tracks possible things with their probabilities
newtype Prob a = Prob { unprob :: [(a, P)] } deriving (Eq, Show)

instance Functor Prob where
  fmap f (Prob xs) = Prob [(f x, p) | (x, p) <- xs]

instance Applicative Prob where
  pure x = Prob [(x, P 1 1)]
  Prob fs <*> Prob xs = Prob [(f x, pf `times` px) | (f, pf) <- fs, (x, px) <- xs]

instance Monad Prob where
  Prob xs >>= f =
    Prob [ (y, p `times` p')
         | (x, p) <- xs
         , (y, p') <- unprob (f x)
         ]

-- probably the wrong name? this is collapsing identical possibilities
-- into one
normalize :: (Ord a) => Prob a -> Prob a
normalize (Prob xs) =
  Prob [(fst (head rs), combine (map snd rs)) | rs <- groupWith fst xs ]