{-# OPTIONS -fglasgow-exts #-}
{-# OPTIONS -O #-}
import qualified Data.Map as Map
import Ix (range)
import Data.Ratio ((%))

-- If you want, skip down to the comment that begins "Now to use this".
-- That's the start of the fun stuff.  All the stuff before that is
-- framework.

data (Ord a, Fractional b) => ProbWorld b a = Probably (Map.Map a b)

instance (Show a, Ord a, Show b, Fractional b) => Show (ProbWorld b a) where
    showsPrec p (Probably x) s = showsPrec p x s

-- >=> is my operator that's like >>=
-- See, Map.fromListWith requires that a implement Ord, so I can't
-- make this into a Monad.  A pity, because it's almost a Monad.
-- Oh well.  I can still use syntax that's almost like "do" by
-- carefully abusing lambda and precedence rules.

(>=>) :: (Ord a, Ord b, Fractional n) => 
            ProbWorld n a -> (a -> ProbWorld n b) -> ProbWorld n b
(Probably m) >=> kn = Probably (Map.fromListWith (+) worlds)
    where
        worlds = [ (k, v) | (ak,av) <- Map.toList m, 
                            (Probably bm) <- [kn ak],
                            (k, v) <- Map.toList (Map.map (* av) bm)]
-- Likewise >>=> is like >>
(>>=>) :: (Ord a, Ord b, Fractional n) => 
            ProbWorld n a -> ProbWorld n b -> ProbWorld n b
(Probably m1) >>=> (Probably m2) = Probably (Map.fromList worlds)
    where
        l1 = Map.elems m1
        l1s = sum l1
        worlds = case l1 of 
                    [] -> []
                    _ -> [ (k, v*l1s) | (k, v) <- Map.toList m2 ]

infixr 0 >=>
infixr 0 >>=>
-- returnP is like return
returnP :: (Ord a, Fractional n) => a -> ProbWorld n a
returnP a = Probably (Map.singleton a (fromInteger 1))

pMap :: (Ord b, Ord c, Fractional n) => 
        (b -> c) -> (ProbWorld n b -> ProbWorld n c)
pMap mapr = pCombFunc
    where
        pCombFunc bp = Probably (Map.mapKeysWith (+) mapr (fromProb bp))
        


fromProb :: (Ord a, Fractional n) => ProbWorld n a -> Map.Map a n
fromProb (Probably m) = m

toList :: (Ord a, Fractional n) => ProbWorld n a -> [(a,n)]
toList p = map (\(k,v) -> (k,v/normv)) l
            where
                normv = foldl (\x y -> x + snd y) 0 l
                l = (Map.toList $ fromProb p)

-- Use this to specify "givens"
assert :: (Fractional n) => Bool -> ProbWorld n ()
assert True  = (returnP ())
assert False = Probably (Map.empty)

-- This means one of the list was chosen at random
choice :: (Ord a, Fractional n) => [a] -> ProbWorld n a
choice [] = Probably (Map.empty)
choice l = Probably $ Map.fromListWith (+) (map toProb l)
    where toProb v = (v, r)
          r = (fromInteger 1) / (fromInteger $ toInteger (length l))

-- Choose one of the "a".  If the Doubles don't add up to "1",
-- your results likely won't mean what you think they do.
choicep :: (Ord a, Fractional n) => [(a, n)] -> ProbWorld n a
choicep [] = Probably (Map.empty)
choicep l = Probably $ Map.fromListWith (+) l

choiceTuple :: (Ord a, Fractional n) => (a,a) -> ProbWorld n a
choiceTuple (x,y) = choice [True, False] >=> \f -> returnP (if f then x else y)

expectedValue :: ProbWorld Double Double -> Double
expectedValue p =
    Map.foldWithKey (\k a b -> b + k*a) (fromInteger 0) (fromProb p)

multichoice :: (Ord a, Fractional n) => Int -> ProbWorld n a -> ProbWorld n [a]
multichoice n c = choiceChain $ replicate n c

choiceChain :: (Ord a, Fractional n) => [ProbWorld n a] -> ProbWorld n [a]
choiceChain [] = returnP []
choiceChain (choice:choices) = choice >=> 
                            \c -> pMap (\cs -> c:cs) (choiceChain choices)

choiceremove :: (Ord a, Fractional n) => [a] -> ProbWorld n (a,[a])
choiceremove [] = Probably (Map.empty)
choiceremove l@(lh:ls) = Probably $ Map.fromListWith (+) mapl
    where 
           llen = length l
           r = (fromInteger 1) / (fromInteger $ toInteger llen)
           mapl = map (\x -> (x, r)) $ take llen chlist
           chlist = map (\(x, fnt, bck) -> (x, myrev fnt bck)) itlist
           itlist = iterate (\(x, fnt, bck) -> (head bck, (x:fnt), tail bck)) (lh, [], ls)
           myrev [] bck = bck
           myrev (x:xs) bck = myrev xs (x:bck)

chooseCombo :: (Ord a, Fractional n) => Int -> [a] -> ProbWorld n [a]
chooseCombo nThings thingList =
        Probably $ Map.fromList [(x, r) | x <- pick nThings thingList]
    where   r = (fromInteger 1) / (fromInteger $ toInteger nComb)
            nComb = pChooseQ (length thingList) nThings

-- A useful formula; we're likely to use it again:
pChooseQ p q
    | p < q  = 0
    | p == q = 1
    | 0 == q = 1
    | otherwise = (p-q+1) * (pChooseQ p (q-1)) `div` q

-- Get all the combinations of n things from a list, but do so in
-- as lazy a manner as possible.
pick :: Int -> [a] -> [[a]]
pick n l = pick' n l 0
  where 
    pick' n lfrom endidx = pick'' n lfrom endidx (pick' n lfrom (endidx+1))
    pick'' n lfrom endidx rest
        | n < 0       = []
        | n == 0      = [[]]
        | n > endidx  = rest
        | otherwise   = pick''' [] n lfrom endidx rest
    pick''' pref 0 _ _ rest = (reverse pref):rest
    pick''' _ _ [] _ _ = []
    pick''' pref n (x:xs) endidx rest
        | n < 0 || n > endidx = rest
        | n == endidx         = pickWithFirst rest
        | n == 1              = pickWithoutFirst rest
        | otherwise           = pickWithFirst (pickWithoutFirst rest)
      where
        pickWithFirst rest'    = pick''' (x:pref) (n-1) xs (endidx-1) rest'
        pickWithoutFirst rest' = pick''' pref n xs (endidx-1) rest'

-- Now to use this.  Let's try some scenarios out.

-- There's a general principle here of "use assert for stuff 
-- that just happens, but use different lists into choice
-- for stuff that was deliberately avoided".  To demonstrate this,
-- first we have the classic "mixed card" problem.

-- Two cards in a pot, one red on both sides and one red on one side,
-- green on the other.  A card is chosen at random and a random
-- side is held up to you.  It is red.  What is the chance that
-- the other side is red?

rgCard :: ProbWorld Rational String
rgCard =    choice ["RR","RG"] >=> \card ->
            choice card >=> \side ->
            assert (side == 'R') >>=>
            returnP ("Other side is " ++ (tail card))

-- Same problem, but the person holding the card up always chooses to
-- show you the red side.  Chances should (obviously) be even now.

rgCard2 :: ProbWorld Rational String
rgCard2 = choice ["RR","RG"] >=> \card ->
          choice (filter (== 'R') card) >=> \side ->
-- This assert is now a no-op; uncommenting it gives the same results
--          assert (side == 'R') >=> \_ ->
          returnP ("Other side is " ++ (tail card))

-- Now let's have some fun...

-- Classic probability problem.
-- Google for "Monty Hall Problem" for description.
-- or, see http://en.wikipedia.org/wiki/Monty_Hall_problem

montyHall :: ProbWorld Rational String
montyHall = choice [1,2,3] >=> \prizedoor ->
            choice [1,2,3] >=> \choicedoor ->
            -- Note the filter below - the host deliberately avoids the prize
            choice (filter (\x -> x /= prizedoor &&
                                x /= choicedoor) [1,2,3]) >=> \hostopen ->
            if (prizedoor == choicedoor)
                 then returnP "You should stay  "
                 else returnP "You should switch"
                 
-- Variation: the host doesn't know where the prize is and opens
-- an unchosen door at random.  If he opens the prize door, the
-- game is reset, and we all try again.
--
-- Now, he just happened to not reveal the prize to you.
-- What should you do?

montyHallDumbHost :: ProbWorld Rational String
montyHallDumbHost = 
                choice [1,2,3] >=> \prizedoor ->
                choice [1,2,3] >=> \choicedoor ->
                -- This choice/filter combo says: The host didn't deliberately
                -- avoid the prize, but just happened to avoid it anyway
                choice (filter (\x -> x /= choicedoor) [1,2,3]) >=> \hostopen ->
                assert (hostopen /= prizedoor) >>=> 
                if (prizedoor == choicedoor)
                  then returnP "You should stay  "
                  else returnP "You should switch"

-- Variation: two players each choose different doors.
-- The host reveals that one of the players chose poorly
-- and sends them packing.  The other player gets a chance to
-- switch if they want.
--
-- Now, he just happened to not eliminate you.
-- What should you do?

montyHall2player :: ProbWorld Rational String
montyHall2player = 
                choice [1,2,3] >=> \prizedoor ->
                choice [1,2,3] >=> \mychoicedoor ->
                -- Bob can't choose the same door as I did, so...
                choice (filter (\x -> x /= mychoicedoor) [1,2,3]) 
                                                    >=> \bobchoicedoor ->
                -- Now the host chooses someone to eliminate from among
                -- the player(s) who chose poorly
                choice (filter (\x -> (snd x) /= prizedoor) 
                        [("me",mychoicedoor),("bob",bobchoicedoor)])
                                                    >=> \elimchoice ->
                -- He just happened to eliminate Bob
                assert ((fst elimchoice) == "bob") >>=>
                if (prizedoor == mychoicedoor)
                  then returnP "You should stay  "
                  else returnP "You should switch"

-- Does it change if Bob (the eliminated player) got
-- to pick his door first?

montyHall2player' :: ProbWorld Rational String
montyHall2player' = 
                choice [1,2,3] >=> \prizedoor ->
                choice [1,2,3] >=> \bobchoicedoor ->
                choice (filter (\x -> x /= bobchoicedoor) [1,2,3]) 
                                                    >=> \mychoicedoor ->
                -- Now the host chooses someone to eliminate from among
                -- the player(s) who chose poorly
                choice (filter (\x -> (snd x) /= prizedoor) 
                        [("me",mychoicedoor),("bob",bobchoicedoor)])
                                                    >=> \elimchoice ->
                -- He just happened to eliminate Bob
                assert ((fst elimchoice) == "bob") >>=>
                if (prizedoor == mychoicedoor)
                  then returnP "You should stay  "
                  else returnP "You should switch"

-- Okay, but what if the host just likes eliminating Bob,
-- and will kick him out early on whenever possible?

montyHallHatesBob :: ProbWorld Rational String
montyHallHatesBob = 
                choice [1,2,3] >=> \prizedoor ->
                choice [1,2,3] >=> \mychoicedoor ->
                -- Bob can't choose the same door as I did, so...
                choice (filter (\x -> x /= mychoicedoor) [1,2,3]) 
                                                    >=> \bobchoicedoor ->
                -- Now the host chooses someone to eliminate from among
                -- the player(s) who chose poorly, but tries to
                -- choose Bob
                choice (if (bobchoicedoor /= prizedoor)
                        then [("bob",bobchoicedoor)]
                        else [("me",mychoicedoor)] ) >=> \elimchoice ->
                -- He "just happened" to eliminate Bob
                assert ((fst elimchoice) == "bob") >>=>
                if (prizedoor == mychoicedoor)
                  then returnP "You should stay  "
                  else returnP "You should switch"

-- Okay, now suppose that both players pick in secret, and then a player
-- is eliminated who chose poorly.  (If such a player exists; if not, we
-- start over)  When they're eliminated, you aren't told anything about
-- what door they picked.  The remaining player can then keep their own
-- door or pick both of the other doors.

montyHallSecretBob :: ProbWorld Rational String
montyHallSecretBob =
                choice [1,2,3] >=> \prizedoor ->
                choice [1,2,3] >=> \mychoicedoor ->
                choice [1,2,3] >=> \bobchoicedoor ->
                -- Now the host chooses someone to eliminate from among
                -- the player(s) who chose poorly
                choice (filter (\x -> (snd x) /= prizedoor) 
                        [("me",mychoicedoor),("bob",bobchoicedoor)])
                                                    >=> \elimchoice ->
                -- He just happened to eliminate Bob
                assert ((fst elimchoice) == "bob") >>=>
                if (prizedoor == mychoicedoor)
                  then returnP "You should keep just your door  "
                  else returnP "You should take both other doors"

montyHallSecretBobHatesBob :: ProbWorld Rational String
montyHallSecretBobHatesBob =
                choice [1,2,3] >=> \prizedoor ->
                choice [1,2,3] >=> \mychoicedoor ->
                choice [1,2,3] >=> \bobchoicedoor ->
                -- Now the host chooses someone to eliminate from among
                -- the player(s) who chose poorly, but tries to
                -- choose Bob
                choice (if (bobchoicedoor /= prizedoor)
                        then [("bob",bobchoicedoor)]
                        else [("me",mychoicedoor)] ) >=> \elimchoice ->
                -- He just happened to eliminate Bob
                assert ((fst elimchoice) == "bob") >>=>
                if (prizedoor == mychoicedoor)
                  then returnP "You should keep just your door  "
                  else returnP "You should take both other doors"

-- You'd never think your buddy would cheat - okay, 
-- maybe there's a 1-in-a-million chance he'd cheat -
-- but he just rolled 10 sixes in a row...

-- cheatChance is the base chance he's cheating (1/1000000)
-- cheatProb is the probability a cheating die rolls a "6"

cheatingDice :: Rational -> Rational -> Int -> ProbWorld Rational String
cheatingDice cheatChance cheatProb nSixes =
    choicep [(True,cheatChance),(False,1-cheatChance)] >=> \isCheater ->
    multichoice nSixes (rollASix isCheater) >>=>
    if (isCheater)
        then returnP "Cheating"
        else returnP "Not Cheating"
    where
        rollASix isCheating = if (isCheating) 
                                then (choicep [(True,cheatProb),(False,1-cheatProb)] >=> \x -> assert x)
                                else (choice [1,2,3,4,5,6] >=> \x -> assert (x==6))


-- Okay, so suppose I choose n bits (0 or 1) at random and add them up.
-- What's the distribution of the sum?

bitsum :: Int -> ProbWorld Rational Int
bitsum n =
    multichoice n (choice [0,1]) >=> \bits ->
    returnP (sum bits)
    
-- Let's play the lottery!
-- except not really, because we run out of heap very quickly
-- I know, very sad.  I suspect I need to be more careful in
-- what I make non-lazy.  I think I could do this with the 
-- proper amount of laziness.  Of course then the computation
-- might still take forever.

-- We use 25 white balls instead of the 55 used in the real thing.

powerballNaive :: ProbWorld Rational String
powerballNaive = 
    chooseCombo 5 (range (1,25)) >=> \whiteBalls ->
    choice (range (1,42)) >=> \redBall ->
    returnP (shows (length (filter (\x -> x < 6) whiteBalls))
        (if (redBall == 1) then " / 1" else " / 0"))

-- Now we can also play it with math, but what fun is that?
-- On the plus side, we get all 55 white balls now
-- We use basic combinatorics to get the number of white
-- balls chosen in common with what we picked.

powerballClever :: ProbWorld Rational String
powerballClever =
    choicep whiteMatchList >=> \whiteMatchCount ->
    choice (range (1,42)) >=> \redBall ->
    returnP (shows whiteMatchCount
             (if (redBall == 1) then " / 1" else " / 0"))
      where
        whiteMatchList =
            [ (x, (pChooseQ 5 x) * (pChooseQ 50 (5-x)) % totChoices) |
              x <- [0,1,2,3,4,5] ]
        totChoices = pChooseQ 55 5


showReport :: (Fractional n, Show n, 
                Show a, Ord a) => String -> ProbWorld n a -> IO ()
showReport header p = do
                putStrLn header
                l <- return (toList p)
                putStr (foldr (\(v,p) s -> "   " ++
                               shows v (" ==> " ++ shows p ('\n':s))) "" l)

toDoubleProb :: (Ord a) => ProbWorld Rational a -> ProbWorld Double a
toDoubleProb = \(Probably m) -> Probably (Map.map fromRational m)

main :: IO ()
main = do
        showReport "Red-Green card problem:" rgCard
        showReport "same with always showing the red side:" rgCard2
        putStrLn ""
        showReport "Basic Monty Hall Scenario:" montyHall
        showReport "Monty Hall with do-overs:" montyHallDumbHost
        showReport "Monty Hall 2 players:" montyHall2player
        showReport "Monty Hall 2 players':" montyHall2player'
        showReport "Monty Hall hates Bob:" montyHallHatesBob
        showReport "Monty Hall secret Bob:" montyHallSecretBob
        showReport "Monty Hall secret Bob, hates Bob:" montyHallSecretBobHatesBob
{-         putStrLn ""
 -         showReport "Bitsum 6:" (bitsum 6)
 -         putStrLn ""
 -         showReport "PowerBall! (Naive, with 25 white balls)" powerballNaive
 -         showReport "PowerBall! (Using math shortcuts, 55 balls)" $ toDoubleProb powerballClever
 -         putStrLn ""
 -         showReport "Cheating 1/1000000 1 10" $ toDoubleProb $ cheatingDice (1%1000000) (1) 10
 -}