module type T = sig
  type 'a t
  val return : 'a -> 'a t
  val bind : 'a t -> ('a -> 'b t) -> 'b t

  val fail : 'a t
  val flip : bool t
  val choice : float -> bool t
end

module P = struct

  type 'a t = ('a -> float) -> float

  let return x = fun f -> f x
  let bind m n = fun f -> m (fun x -> (n x) f)

  let fail = fun f -> 0.
  let choice p = fun f -> p *. f true +. (1.-.p) *. f false
  let flip = choice 0.5

end

module Algos (M:T) = struct

  let pair =
    M.bind M.flip (fun x -> M.bind M.flip (fun y -> M.return (x,y)))

  let pick l =
    let rec aux l len =
      match l with
        | [] -> M.fail
        | hd::tl ->
            M.bind
              (M.choice (1. /. float len))
              (fun b ->
                 if b then
                   M.return hd
                 else
                   aux tl (len-1))
    in
      aux l (List.length l)

end

let dirac x = fun y -> if x = y then 1. else 0.

module TestP = struct
  module A = Algos(P)
  let () =
    Printf.printf "P[pair=(true,true)] = %.2f\n" (A.pair (dirac (true,true))) ;
    Printf.printf "P[pick([1;2;3])=2] = %.3f\n" (A.pick [1;2;3] (dirac 2)) ;
    Printf.printf "P[pick([1;2;2])=2] = %.3f\n" (A.pick [1;2;2] (dirac 2))
end

module Chevre = struct

  module A = Algos(P)
  let (>>=) = P.bind

  let doors = [1;2;3]

  type second = int -> int P.t
  type first = (int*second) P.t

  let play (strategy:first) =
    A.pick doors >>= fun tresor ->
    strategy >>= fun (first,strategy) ->
    let empty_doors =
      List.filter
        (fun d -> d <> first && d <> tresor)
        doors
    in
    A.pick empty_doors >>= fun empty ->
    strategy empty >>= fun second ->
    P.return (second = tresor)

  let () =
    let second : second =
      fun _ -> A.pick doors
    in
    let first : first =
      P.bind
        (A.pick doors)
        (fun p -> P.return (p, second))
    in
      Printf.printf "P[win] = %.3f\n" (play first (dirac true))

  let () =
    let second first : second =
      fun free ->
        P.return (List.find (fun p -> p <> free && p <> first) doors)
    in
    let first : first =
      P.bind
        (A.pick doors)
        (fun p -> P.return (p, second p))
    in
      Printf.printf "P[win] = %.3f\n" (play first (dirac true))

end