;; Information Theory

(defun lg (x)
  (log x 2))

(defun information (p)
  "Return the amount of information given probability P in bits."
  (if (zerop p)
      single-float-positive-infinity
      (lg (/ p))))

(defun entropy (q)
  "Return the expected value of information in messages of Q
probabilities."
  (reduce #'+ q :key (lambda (p) (* p (information p)))))

(defun cross-entropy (p x y)
  "Return a measure of mutual information (dependency) between two
random variables X and Y.

If the variables are independent, the mutual information between them
is zero.

P should be a function that takes two values X[i] and Y[i] and returns
the probability of both being true.  NIL may be passed for either of
them, in which case the probability returned should be for the other
value being true."
  (flet ((inner (xi yi)
           (let ((p-x (funcall p xi nil)))
             (if (zerop p-x)
                 single-float-positive-infinity
                 (let ((p-y (funcall p nil yi)))
                   (if (zerop p-y)
                       single-float-positive-infinity
                       (let ((p-xy (funcall p xi yi)))
                         (if (zerop p-xy)
                             0.0
                             ;; Eq. 2.24 in the book is missing
                             ;; parentheses for the denominator.
                             (* p-xy (lg (/ p-xy (* p-x p-y))))))))))))
    (reduce #'+ x :key (lambda (xi) (reduce #'+ y :key (lambda (yi) (inner xi yi)))))))

(defun conditional-mutual-information (p x y z)
  "Return the mutual information (dependency) between two random
variables X and Y, given random variable Z.

P should be a function that takes three values X[i], Y[i], and Z[i]
and returns the following probabilities:

  (P NIL NIL Z) -> P(Z)

  (P X NIL Z) -> P(X|Z)

  (P NIL Y Z) -> P(Y|Z)

  (P X Y Z) -> P(X,Y|Z)"
  (flet ((inner (xi yi zi)
           (let ((p-xz (funcall p xi nil zi)))
             (if (zerop p-xz)
                 single-float-positive-infinity
                 (let ((p-yz (funcall p nil yi zi)))
                   (if (zerop p-yz)
                       single-float-positive-infinity
                       (let ((p-xyz (funcall p xi yi zi)))
                         (if (zerop p-xyz)
                             0.0
                             (* p-xyz
                                (lg (/ p-xyz (* p-xz p-yz))))))))))))
    (reduce #'+ z :key (lambda (zi)
                         (* (funcall p nil nil zi)
                            (reduce #'+ y :key (lambda (yi) (reduce #'+ x :key (lambda (xi) (inner xi yi zi))))))))))