;;; PLT Scheme Science Collection
;;; random-distributions/poisson.ss
;;; Copyright (c) 2004-2007 M. Douglas Williams
;;;
;;; This library is free software; you can redistribute it and/or
;;; modify it under the terms of the GNU Lesser General Public
;;; License as published by the Free Software Foundation; either
;;; version 2.1 of the License, or (at your option) any later version.
;;;
;;; This library is distributed in the hope that it will be useful,
;;; but WITHOUT ANY WARRANTY; without even the implied warranty of
;;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
;;; Lesser General Public License for more details.
;;;
;;; You should have received a copy of the GNU Lesser General Public
;;; License along with this library; if not, write to the Free
;;; Software Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA
;;; 02111-1307 USA.
;;;
;;; -------------------------------------------------------------------
;;;
;;; This module implements the poisson distribution.  It is based on
;;; the Random Number Distributions in the GNU Scientific Library.
;;;
;;; Version  Date      Description
;;; 1.0.0    09/28/04  Marked as ready for Release 1.0.  Added
;;;                    contracts for functions (Doug Williams)
;;; 2.0.0    11/19/07  Added unchecked versions of functions and
;;;                    getting ready for PLT Scheme 4.0 (Doug Williams)

(module poisson mzscheme
  
  (require (lib "contract.ss"))
  
  (provide
   (rename random-poisson unchecked-random-poisson)
   (rename poisson-pdf unchecked-poisson-pdf))
  
  (provide/contract
   (random-poisson
    (case-> (-> random-source? real? natural-number/c)
            (-> real? natural-number/c)))
   (poisson-pdf
    (-> natural-number/c real? (>=/c 0.0))))
  
  (require "../random-source.ss")
  (require "../special-functions/gamma.ss")
  (require "gamma.ss")
  (require "binomial.ss")
  
  ;; random-poisson: random-source x real -> integer
  ;; random-poisson: real -> integer
  ;;
  ;; This function returns a random variate from a poisson distribution
  ;; with mean mu.
  (define random-poisson
    (case-lambda
      ((r mu)
       (let/ec exit
         (let ((k 0))
           (let loop ()
             (if (> mu 10.0)
                 (let* ((m (inexact->exact
                            (truncate (* mu (/ 7.0 8.0)))))
                        (x (unchecked-random-gamma-int r m)))
                   (if (>= x mu)
                       (exit (+ k (unchecked-random-binomial r (/ mu x) (- m 1))))
                       (begin
                         (set! k (+ k m))
                         (set! mu (- mu x))
                         (loop))))))
           ;; This following method works well when mu is small
           (let ((emu (exp (- mu)))
                 (prod 1.0))
             (let loop ()
               (set! prod (* prod (unchecked-random-uniform r)))
               (set! k (+ k 1))
               (if (> prod emu)
                   (loop)))
             (- k 1)))))
      ((mu)
       (random-poisson (current-random-source) mu))))
  
  ;; poisson-pdf: integer -> real
  ;;
  ;; This function computes the probability density p(x) at x for a
  ;; poisson distribution with mean mu.
  (define (poisson-pdf k mu)
    (let ((lf (unchecked-lnfact k)))
      (exp (- (* (log mu) k) lf mu))))
  
)