(use 'clojure.contrib.monads)
(use 'clojure.contrib.probabilities.finite-distributions)

(defn inc-m [x]
  (when (not= x 6)
    (inc x)))

(assert (= [1 2 3 4 5 6 nil 8 9 10]
           (map inc-m (range 10))))

(defn double-m [x]
  (when (even? x)
    (* 2 x)))

(assert (= [0 nil 4 nil 8 nil 12 nil 16 nil]
           (map double-m (range 10))))

(defn dec-m [x]
  (when (not= x 3)
    (dec x)))

(assert (= [-1 0 1 nil 3 4 5 6 7 8]
           (map dec-m (range 10))))

(defn all-m [x]
  (when-not (nil? x)
    (let [x1 (inc-m x)]
      (when-not (nil? x1)
        (let [x2 (double-m x1)]
          (when-not (nil? x2)
            (dec-m x2)))))))

(assert (= [nil 3 nil 7 nil 11 nil 15 nil 19]
           (map all-m (range 10))))

(with-monad maybe-m
            (def all-m (m-chain [inc-m
                               double-m
                               dec-m])))

(assert (= [nil 3 nil 7 nil 11 nil 15 nil 19]
           (map all-m (range 10))))


; the probability monad
  
(defn inc-p [x]
  {x 1/2 (inc x) 1/2})

(assert (= {5 1/2 6 1/2}
           (inc-p 5)))

(defn double-p [x]
  {x 3/4 (* 2 x) 1/4})

(assert (= {3 3/4 6 1/4}
           (double-p 3)))

(defn all-p [x]
  (reduce (fn [dist d]
            (merge-with + dist d)) 
          (for [[y1 p1] (inc-p x)
                [y2 p2] (double-p y1)]
            {y2 (* p1 p2)})))

(assert (= {8 1/8, 4 3/8, 6 1/8, 3 3/8}
           (all-p 3)))

(with-monad dist-m
            (def all-p (m-chain [inc-p
                                 double-p])))

(assert (= {8 1/8, 4 3/8, 6 1/8, 3 3/8}
           (all-p 3)))