Prisoners.hs

import Data.List.Extras.Argmax

data Choice = PickRed | PickGreen | Abstain deriving (Show,Eq)
data Results = PilloriedInScience | EatenByTigers | Freed deriving (Eq,Show)
data Color = Red | Green deriving (Show,Eq)

-- to simplify things, let's assume choices are only predicated upon
-- the number of bits that go one way or the other
-- (numRed, numGreen)
-- Note that ChoiceFunctions are a function of the SEEN input, so the sum
-- of the numRed and numGreen that they get should be one less than the total.
type BoothInput = (Int,Int)
type ChoiceFunction  = BoothInput -> Choice

generateAllInputs :: Int -> [[Color]]
generateAllInputs 1 = [[Red],[Green]]
generateAllInputs n = let rec = generateAllInputs (n-1) in
  (map (\x-> Red:x) rec) ++ (map (\x-> Green:x) rec)

generateAllChoiceLists :: Int -> [[Choice]]
generateAllChoiceLists 1 = [[PickRed],[PickGreen],[Abstain]]
generateAllChoiceLists n = let rec = generateAllChoiceLists (n-1) in
  (map (\x->PickRed:x) rec) ++ (map (\x-> PickGreen:x) rec) ++ (map (\x-> Abstain:x) rec)

generateAllChoiceFuncs :: Int -> [ChoiceFunction]
generateAllChoiceFuncs 0 = []
generateAllChoiceFuncs n =
  let
    allChoiceLists = generateAllChoiceLists n
    inputs = map (\x-> (x,n-1-x)) [0..(n-1)]
  in
    map (graphToFunc . (zip inputs)) allChoiceLists

-- very unsafe and unoptimized, but whatevs. yay lazy languages
graphToFunc :: (Eq a, Show a) => [(a,b)] -> (a->b)
graphToFunc [] = \x -> error ("Didn't handle input: " ++ (show x))
graphToFunc ((x,fx):tail) = \y -> if y == x then fx else graphToFunc tail y

showChoiceFunction :: ChoiceFunction -> Int -> String
showChoiceFunction f n =
  let
    domain = map (\x->(x,n-1-x)) [0..(n-1)]
    itemDec = map (\(x,y) -> "\t("++(show x)++" Red,"++(show y)++" Green) -> "++(show$f(x,y))++"\n") domain
  in
    (foldl (++) "{\n" itemDec) ++ "}\n"

choiceFuncGoodness :: ChoiceFunction -> [[Color]] -> Float
choiceFuncGoodness func inputs =
  let
    goodResults = filter (\x->x==Freed) (map (getResults func . numRedGreen) inputs)
  in
    (fromIntegral $ length goodResults) / (fromIntegral $ length inputs)

getBestFunc :: [ChoiceFunction] -> [[Color]] -> ChoiceFunction
getBestFunc funcs inputs = argmax (\x -> choiceFuncGoodness x inputs) funcs

getResults :: ChoiceFunction -> BoothInput -> Results
getResults choiceFunc input =
  endResult $
       (replicate (fst input) $ (choiceFunc (seenInput Red input), Red))
    ++ (replicate (snd input) $ (choiceFunc (seenInput Green input), Green))

numRedGreen :: [Color] -> BoothInput
numRedGreen [] = (0,0)
numRedGreen (Red:tail) = let rec = numRedGreen tail in (1+fst rec, snd rec)
numRedGreen (Green:tail) = let rec = numRedGreen tail in (fst rec, 1+snd rec)


seenInput :: Color -> BoothInput -> BoothInput
seenInput Red (x,y) = (x-1,y)
seenInput Green (x,y) = (x,y-1)

endResult :: [(Choice,Color)] -> Results
endResult [] = EatenByTigers
endResult ((Abstain,_):tail) = endResult tail
endResult ((PickRed,Red):tail) = case endResult tail of
  PilloriedInScience -> PilloriedInScience
  _ -> Freed
endResult ((PickGreen,Green):tail) = endResult ((PickRed,Red):tail)
endResult _ = PilloriedInScience

bestFuncForSize :: Int -> ChoiceFunction
bestFuncForSize n = getBestFunc (generateAllChoiceFuncs n) (generateAllInputs n)

showFuncWithGoodness :: Int -> ChoiceFunction -> String
showFuncWithGoodness n f =
  let
    funcRep = showChoiceFunction f n
    funcVal = choiceFuncGoodness f (generateAllInputs n)
    output = funcRep ++ "\nChance of escape: " ++ (show funcVal) ++ "\n"
  in
    output

main =
  let
    numBooths = 5
    bestFunc = bestFuncForSize numBooths
  in
    putStr $ showFuncWithGoodness numBooths bestFunc
    --putStr $ foldl (++) "" $ map (showFuncWithGoodness numBooths) (generateAllChoiceFuncs numBooths)