(* The Caml code for the various monads of lecture 4 *)

module type MONAD = sig
  type 'a mon
  val ret: 'a -> 'a mon
  val bind: 'a mon -> ('a -> 'b mon) -> 'b mon
  type 'a result
  val run: 'a mon -> 'a result
end

(* The Exception monad, a.k.a. the Error monad *)

module Exception = struct

  type 'a mon = V of 'a | E of exn

  let ret (x: 'a) : 'a mon = V x
  let bind (m: 'a mon) (f: 'a -> 'b mon): 'b mon =
    match m with E exn -> E exn | V x -> f x

  type 'a result = 'a 
  let run (m: 'a mon) : 'a result =
    match m with V x -> x | E exn -> failwith "uncaught exception"

  let raise (x: exn) : 'a mon = E x
  let trywith (m: 'a mon) (f: exn -> 'a mon): 'a mon =
    match m with E exn -> f exn | V x -> V x
end

(* The State monad, specialized to references over integers *)

module Store = struct
  module IntMap = Map.Make(struct type t = int let compare = compare end)
  type t = { next: int; contents: int IntMap.t }
  type ref = int
  let empty = { next = 0; contents = IntMap.empty }
  let alloc x s =
      (s.next,
       { next = s.next + 1; contents = IntMap.add s.next x s.contents })
  let read r s = IntMap.find r s.contents
  let write r x s =
     { next = s.next; contents = IntMap.add r x s.contents }
end

module State = struct
  type 'a mon = Store.t -> 'a * Store.t

  let ret (x: 'a) : 'a mon = fun v -> (x, v)
  let bind (x: 'a mon) (f: 'a -> 'b mon) : 'b mon =
    fun v -> let (x', v') = x v in f x' v'

  type 'a result = 'a 
  let run (c: 'a mon) : 'a result = fst (c Store.empty)

  let ref (n : int) : Store.ref mon =
    fun s -> Store.alloc n s
  let deref (r: Store.ref) : int mon =
    fun s -> (Store.read r s, s)
  let assign (r: Store.ref) (x: int) : unit mon =
    fun s -> ((), Store.write r x s)
end

let sumlist l =
  State.run (State.bind (State.ref 0) (fun r ->
    let rec sum = function
    | [] -> State.deref r
    | hd :: tl -> State.bind (State.deref r) (fun n ->
                  State.bind (State.assign r (n + hd)) (fun _ ->
                  sum tl))
    in sum l))

(* The Continuation monad, assuming the type of answers is "int" *)

module Cont = struct

  type answer = int
  type 'a mon = ('a -> answer) -> answer

  let ret (x: 'a) : 'a mon = fun k -> k x
  let bind (m: 'a mon) (f: 'a -> 'b mon) : 'b mon =
    fun k -> m (fun x -> f x k)

  let run (m: answer mon) = m (fun x -> x)

  let callcc (f: ('a -> answer) -> 'a mon) : 'a mon =
    fun k -> f k k
  let throw (k: 'a -> answer) (x: 'a) : 'a mon =
    fun _ -> k x
end

let exple_cont n =
  Cont.run 
    (Cont.callcc (fun k ->
       Cont.bind
         (if n < 0 then Cont.throw k n else Cont.ret n)
         (fun x -> Cont.ret (x + 1))))

(* The Logging monad *)

module Log = struct
  type log = string list
  type 'a mon = log -> 'a * log

  let ret (x: 'a) : 'a mon = fun l -> (x, l)
  let bind (m: 'a mon) (f: 'a -> 'b mon): 'b mon =
    fun l -> match m l with (x, l') -> f x l'

  type 'a result = 'a * log
  let run (m: 'a mon): 'a result =
    match m [] with (x, l) -> (x, List.rev l)

  let log msg : unit mon = fun l -> ((), msg :: l)
end

let log_abs n =
  Log.run
    (if n >= 0 
     then Log.bind (Log.log "positive") (fun _ -> Log.ret n)
     else Log.bind (Log.log "negative") (fun _ -> Log.ret (-n)))

(* Non-determinism, a.k.a. the List monad *)

module Nondet = struct

  type 'a mon = 'a list

  let ret (x: 'a): 'a mon = [x]
  let rec bind (m: 'a mon) (f: 'a -> 'b mon): 'b mon =
    match m with [] -> [] | hd :: tl -> f hd @ bind tl f

  type 'a result = 'a list
  let run (m: 'a mon): 'a result = m

  let fail : 'a mon = []
  let either (a: 'a mon) (b: 'a mon): 'a mon = a @ b
end

let rec insert x l =
  Nondet.either (Nondet.ret (x :: l))
    (match l with
     | [] -> Nondet.fail
     | hd :: tl ->
         Nondet.bind (insert x tl)
                     (fun l' -> Nondet.ret (hd :: l')))

let rec permut l =
  match l with
  | [] -> Nondet.ret []
  | hd :: tl ->
      Nondet.bind (permut tl) (fun l' -> insert hd l')

let test_nondet =
  Nondet.run (permut [1;2;3])

(* Randomized computations *)

module type RANDOM = sig
  type 'a mon
  val ret: 'a -> 'a mon
  val bind: 'a mon -> ('a -> 'b mon) -> 'b mon
  val choose: float -> 'a mon -> 'a mon -> 'a mon
  val rand: int -> int mon
end

module Random_examples(M: RANDOM) = struct

  let roll_3d6 =
    M.bind (M.rand 6) (fun d1 ->
    M.bind (M.rand 6) (fun d2 ->
    M.bind (M.rand 6) (fun d3 -> 
        M.ret (1+d1 + 1+d2 + 1+d3))))

  type light = Green | Yellow | Red

  let traffic_light =
    M.choose 0.1 (M.ret Yellow)
                 (M.choose 0.5 (M.ret Red) 
                               (M.ret Green))
end

(* The Simulation monad *)

module Simulation = struct
  type state = int
  type 'a mon = state -> 'a * state

  let ret (x: 'a) : 'a mon = fun s -> (x, s)
  let bind (m: 'a mon) (f: 'a -> 'b mon) : 'b mon =
    fun s -> let (x', s') = m s in f x' s'
  let run (seed: state) (m: 'a mon) = fst (m seed)

  let next_state (s: state) : state = s * 25173 + 1725

  let rand (n: int) : int mon =
    fun s -> ((abs s) mod n, next_state s)

  let choose (p: float) (a: 'a mon) (b: 'a mon) : 'a mon =
    fun s ->
      if float (abs s) <= p *. float max_int
      then a (next_state s)
      else b (next_state s)
end

module S = Random_examples(Simulation)

let exple1 = Simulation.run 123456 S.roll_3d6
let exple2 = Simulation.run 7890 S.roll_3d6
let exple3 = Simulation.run 4567 S.traffic_light

(* The Distribution monad *)

module Distribution = struct

  type 'a mon = ('a * float) list (* sum of floats = 1.0 *)

  let ret (x: 'a) : 'a mon = [x, 1.0]
  let mulp p l = List.map (fun (x, p') -> (x, p *. p')) l
  let rec bind (x: 'a mon) (f: 'a -> 'b mon) : 'b mon =
    match x with [] -> [] | (hd, p) :: tl -> mulp p (f hd) @ bind tl f

  let choose (p: float) (x: 'a mon) (y: 'a mon) : 'a mon =
    mulp p x @ mulp (1.0 -. p) y
  let rand (n: int) : int mon =
    let p = 1.0 /. float n in
    let rec rnd n = if n < 0 then [] else (n, p) :: rnd (n-1) in
    rnd (n-1)

  let rec accumulate x = function
    | [] -> 0.0
    | (hd, p) :: tl -> if x = hd then p +. accumulate x tl else accumulate x tl

  let rec filter x = function
    | [] -> []
    | (hd, p as p_hd) :: tl -> if x = hd then filter x tl else p_hd :: filter x tl

  let rec compact (x: 'a mon) : 'a mon =
    match x with
    | [] -> x
    | [_] -> x
    | (hd, p) :: tl -> (hd, accumulate hd x) :: compact (filter hd tl)

  let run (m: 'a mon) = compact m
end

module D = Random_examples(Distribution)

let exple5 = Distribution.run D.roll_3d6
let exple6 = Distribution.run D.traffic_light


(* The Expectation monad *)

module Expectation = struct

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

  let ret (x: 'a) : 'a mon = fun k -> k x
  let bind (m: 'a mon) (f: 'a -> 'b mon) : 'b mon =
    fun k -> m (fun x -> f x k)

  let choose (p: float) (x: 'a mon) (y: 'a mon) : 'a mon =
    fun k -> p *. x k +. (1.0 -. p) *. y k

  let rand (n: int) : int mon =
    fun k -> 
      let rec sum n = if n <= 0 then 0.0 else k (n-1) +. sum (n-1) in
      sum n /. float n

  (* Return the expectation of a given result value *)
  let run (res: 'a) (m: 'a mon) =
    m (fun x -> if x = res then 1.0 else 0.0)
end

module E = Random_examples(Expectation)

let exple8 = Expectation.run 16 E.roll_3d6
let exple9 = Expectation.run E.Green E.traffic_light


(**** Monad transformers *)

(* The Identity monad *)

module Identity = struct
  type 'a mon = 'a
  let ret x = x
  let bind m f = f m
  type 'a result = 'a
  let run m = m
end

(* The monad transformer for exceptions *)

module ExceptionTransf(M: MONAD) = struct
  type 'a outcome = V of 'a | E of exn
  type 'a mon = ('a outcome) M.mon

  let ret x = M.ret (V x)
  let bind m f = 
    M.bind m (function E e -> M.ret (E e) | V v -> f v)
  let lift x = M.bind x (fun v -> M.ret (V v))

  type 'a result = 'a M.result

  let run m = M.run (M.bind m (function V x -> M.ret x))

  let raise e = M.ret (E e)
  let trywith m f =
    M.bind m (function E e -> f e | V v -> M.ret (V v))
end

(* The monad transformer for state *)

module StateTransf(M: MONAD) = struct

  type 'a mon = Store.t -> ('a * Store.t) M.mon

  let ret x = fun s -> M.ret (x, s)
  let bind m f =
    fun s -> M.bind (m s) (fun (x, s') -> f x s')
  let lift m = fun s -> M.bind m (fun x -> M.ret (x, s))

  type 'a result = 'a M.result

  let run m =
    M.run (M.bind (m Store.empty) (fun (x, s') -> M.ret x))

  let ref x = fun s -> M.ret (Store.alloc x s)
  let deref r = fun s -> M.ret (Store.read r s, s)
  let assign r x = fun s -> M.ret (Store.write r x s)
end

(* The monad transformer for continuations *)

module ContTransf(M: MONAD) = struct
  type answer = int
  type 'a mon = ('a -> answer M.mon) -> answer M.mon

  let ret x = fun k -> k x
  let bind m f = fun k -> m (fun v -> f v k)
  let lift m = fun k -> M.bind m k
  let run m = M.run (m (fun x -> M.ret x))

  let callcc f = fun k -> f k k
  let throw c x = fun k -> c x
end

(* Example of combinations *)

module StateAndException = struct
  include ExceptionTransf(State)
  let ref x = lift (State.ref x)
  let deref r = lift (State.deref r)
  let assign r x = lift (State.assign r x)
end

module ContinuationAndState = struct
  include ContTransf(State)
  let ref x = lift (State.ref x)
  let deref r = lift (State.deref r)
  let assign r x = lift (State.assign r x)
end
  
(* The Concurrency monad transformer *)

module Concur(M: MONAD) = struct
  type answer = 
    | Seq of answer M.mon
    | Par of answer * answer
    | Stop
  type 'a mon = ('a -> answer) -> answer

  let ret (x: 'a): 'a mon = fun k -> k x
  let bind (m: 'a mon) (f: 'a -> 'b mon): 'b mon =
    fun k -> m (fun v -> f v k)

  let atom (m: 'a M.mon): 'a mon =
    fun k -> Seq(M.bind m (fun v -> M.ret (k v)))
  let stop : 'a mon = fun k -> Stop
  let par (m1: 'a mon) (m2: 'a mon): 'a mon = fun k -> Par (m1 k, m2 k)

  let rec schedule acts =
    match acts with
    | [] -> M.ret ()
    | Seq m :: rem ->
          M.bind m (fun m' -> schedule (rem @ [m']))
    | Par(a1, a2) :: rem ->
          schedule (a1 :: a2 :: rem)
    | Stop :: rem ->
          schedule rem
  type 'a result = 'a M.result
  let run (m: 'a mon) : 'a result = M.run (schedule [m (fun _ -> Stop)])

end

module M = Concur(Log)

let rec loop n s =
    if n <= 0
    then M.ret ()
    else M.bind (M.atom (Log.log s)) (fun _ -> loop (n-1) s)

M.run (M.bind (M.atom (Log.log "start:")) (fun _ ->
               M.par (loop 6 "a") (loop 4 "b")))