;;; PLT Scheme Science Collection
;;; random-distributions/discrete.ss
;;; Copyright (c) 2004 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 general discrete distribution.
;;;
;;; Version  Date      Description
;;; 1.0.0    09/28/04  Marked as ready for Release 1.0.  Added
;;;                    contracts for functions.  (Doug Williams)

(module discrete mzscheme
  
  ;; Data Definition
  
  ;; discrete structure
  ;;
  ;; An instance of the discrete structure, created by make-discrete,
  ;; represents a general discrete distribution.  Only the discrete?
  ;; function is exported.
  (define-values (struct:discrete
                  discrete-constructor
                  discrete?
                  discrete-field-ref
                  set-discrete-field!)
    (make-struct-type 'discrete #f 2 0))
  
  ;; COntracts
  
  (require (lib "contract.ss"))
  
  (provide/contract
   (discrete?
    (-> any/c boolean?))
   (make-discrete
    (-> (vectorof real?) discrete?))
   (random-discrete
    (case-> (-> random-source? discrete? natural-number/c)
            (-> discrete? natural-number/c)))
   (discrete-pdf
    (-> discrete? integer? (real-in 0.0 1.0)))
   (discrete-cdf
    (-> discrete? integer? (real-in 0.0 1.0))))
  
  (require "../random-source.ss")
  
  (define discrete-p
    (make-struct-field-accessor discrete-field-ref 0 'p))
  
  (define discrete-c
    (make-struct-field-accessor discrete-field-ref 1 'c))
  
  ;; make-discrete: vector -> discrete
  ;;
  ;; This function accepts a vector of weights and returns a discrete
  ;; probability structure that can be passed to random-discrete to
  ;; generate random variates.  The weights do not have to sum to 1.0.
  (define (make-discrete w)
    (let* ((n (vector-length w))
           (p (make-vector n))
           (c (make-vector (+ n 1)))
           (sum 0.0)
           (cumm 0.0))
      ;; find sum
      (do ((i 0 (+ i 1)))
          ((= i n) (void))
        (let ((wi (vector-ref w i)))
          (if (> wi 0.0)
              (set! sum (+ sum wi)))))
      ;; build pdf and cdf
      (do ((i 0 (+ i 1)))
          ((= i n) (void))
        (let* ((wi (vector-ref w i))
               (q (if (< wi 0.0) 0.0 (/ wi sum))))
          (vector-set! p i q)
          (vector-set! c i cumm)
          (set! cumm (+ cumm q))))
      (vector-set! c n 1.0)
      ;; make and return the discrete
      (discrete-constructor p c)))
  
  ;; random-discrete: random-source x discrete -> integer
  ;; random-discrete: discrete -> integer
  ;;
  ;; This function returns a random variate from the given discrete
  ;; distribution given by d.
  (define random-discrete
    (case-lambda
      ((r d)
       (let* ((c (discrete-c d))
              (n (- (vector-length c) 1))
              (u (random-uniform r)))
         (let loop ((i 0))
           (if (<= u (vector-ref c (+ i 1)))
               i
               (loop (+ i 1))))))
      ((d)
       (random-discrete (current-random-source) d))))
  
  ;; discrete-pdf: discrete x integer -> real
  ;;
  ;; This function computes the probability density p(k) at k for a
  ;; discrete distribution given by d.
  (define (discrete-pdf d k)
    (let* ((p (discrete-p d))
           (n (vector-length p)))
      (if (or (< k 0)
              (>= k n))
          0.0
          (vector-ref p k))))
  
  ;; discrete-cdf: discrete x integer -> real
  ;;
  ;; This function computes the cummulative density d(k) at k for a
  ;; discrete distribution given by d.
  (define (discrete-cdf d k)
    (let* ((c (discrete-c d))
           (n (- (vector-length c) 1)))
      (cond ((< k 0)
             0.0)
            ((>= k n)
             1.0)
            (else
             (vector-ref c (+ k 1))))))
  
)