#lang racket
;;; Racket Simulation Collection
;;; history.rkt
;;; Copyright (c) 2005-2010 M. Douglas Williams
;;;
;;; This file is part of the Racket Simulation Collection.
;;;
;;; The Racket Simulation Collection 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 3 of the License,
;;; or (at your option) any later version.
;;;
;;; The Racket Simulation Collection 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 the Racket Simulation Collection.  If not, see
;;; <http://www.gnu.org/licenses/>.
;;;
;;; -----------------------------------------------------------------------------
;;;
;;; Version  Date      Description
;;; 4.0.0    09/01/11  Initial release. (MDW)

(require (planet williams/science/random-source)
         (planet williams/science/random-distributions-with-graphics))

;;; random-distribution-hash : hash-eq?
;;; An eq? hash table mapping names to random distributions.
(define random-distribution-hash (make-hasheq))

;;; (find-random-distribution sym) -> random-distribution?
;;;   sym : symbol?
(define (find-random-distribution sym)
  (hash-ref random-distribution-hash sym
            (lambda ()
              (error
               (format "~a is not the name of a random distribution" sym)))))

;;; (struct random-distribution (spec
;;;                              form
;;;                              sample
;;;                              pdf
;;;                              cdf
;;;                              plot)
;;;   spec : (cons/c symbol (listof real?))
;;;   form : string?
;;;   sample : procedure?
;;;   pdf : procedure?
;;;   cdf : procedure?
;;;   plot : procedure?
(struct random-distribution
  (spec
   form
   sample
   pdf
   cdf
   plot))

;;; (define-random-distribution (name . args)
;;;   form random pdf cdf plot)
;;; Defines name to be a random distribution.
(define-syntax-rule (define-random-distribution (name . args)
                      form random pdf cdf plot)
  (hash-set! random-distribution-hash
             'name
             (random-distribution
              '(name . args)
              form
              random
              pdf
              cdf
              plot)))

;;; Random Distributions

;;; Continuous Distributions

(define-random-distribution (beta a b)
  "beta distribution with parameters ~a and ~a"
  random-beta
  beta-pdf
  beta-cdf
  beta-plot)

(define-random-distribution (chi-squared nu)
  "chi squared distribution with ~a degrees of freedom"
  random-chi-squared
  chi-squared-pdf
  chi-squared-cdf
  chi-squared-plot)

(define-random-distribution (exponential mu)
  "exponential distribution with mean ~a"
  random-exponential
  exponential-pdf
  exponential-cdf
  exponential-plot)

(define-random-distribution (f-distribution nu1 nu2)
  "F-distribution with ~a and ~a degrees of freedom"
  random-f-distribution
  f-distribution-pdf
  f-distribution-cdf
  f-distribution-plot)

(define-random-distribution (flat a b)
  "flat (uniform) distribution from ~a to ~a"
  random-flat
  flat-pdf
  flat-cdf
  flat-plot)

(define-random-distribution (gamma a b)
  "gamma distribution with parameters ~a and ~a"
  random-gamma
  gamma-pdf
  gamma-cdf
  gamma-plot)

(define-random-distribution (gaussian mu sigma)
  "Gaussian (normal) distribution with mean ~a and standard deviation ~a"
  random-gaussian
  gaussian-pdf
  gaussian-cdf
  gaussian-plot)

(define-random-distribution (gaussian-tail a mu sigma)
  "upper tail of the Gaussian (normal) distribution with lower limit ~a, mean ~a and standard deviation ~a"
  random-gaussian-tail
  gaussian-tail-pdf
  #f
  gaussian-tail-plot)

(define-random-distribution (lognormal mu sigma)
  "log normal distribution with mean ~a and standard deviation ~a"
  random-lognormal
  lognormal-pdf
  lognormal-cdf
  lognormal-plot)

(define-random-distribution (normal mu sigma)
  "Gaussian (normal) distribution with mean ~a and standard deviation ~a"
  random-gaussian
  gaussian-pdf
  gaussian-cdf
  gaussian-plot)

(define-random-distribution (pareto a b)
  "Pareto distribution with parameters ~a and ~a"
  random-pareto
  pareto-pdf
  pareto-cdf
  pareto-plot)

(define-random-distribution (standard-normal)
  "Gaussian (normal) distribution with mean 0.0 and standard deviation 1.0"
  random-unit-gaussian
  unit-gaussian-pdf
  unit-gaussian-cdf
  unit-gaussian-plot)

(define-random-distribution (standard-normal-tail a)
  "upper tail of the Gaussian (normal) distribution with lower limit ~a, mean 0.0 and standard deviation 1.0"
  random-unit-gaussian-tail
  unit-gaussian-tail-pdf
  #f
  unit-gaussian-tail-plot)

(define-random-distribution (t-distribution nu)
  "t-distribution with ~a degrees of freedom"
  random-t-distribution
  t-distribution-pdf
  t-distribution-cdf
  t-distribution-plot)

(define-random-distribution (triangular a b c)
  "triangular distribution with minimum value ~a, maximum value ~a, and most likely value ~a"
  random-triangular
  triangular-pdf
  triangular-cdf
  triangular-plot)

(define-random-distribution (unit-gaussian)
  "Gaussian (normal) distribution with mean 0.0 and standard deviation 1.0"
  random-unit-gaussian
  unit-gaussian-pdf
  unit-gaussian-cdf
  unit-gaussian-plot)

(define-random-distribution (unit-gaussian-tail a)
  "upper tail of the Gaussian (normal) distribution with lower limit ~a, mean 0.0 and standard deviation 1.0"
  random-unit-gaussian-tail
  unit-gaussian-tail-pdf
  #f
  unit-gaussian-tail-plot)

;;; Discrete Distributions

(define-random-distribution (bernoulli p)
  "Bernoulli distribution with probability ~a"
  random-bernoulli
  bernoulli-pdf
  bernoulli-cdf
  bernoulli-plot)

(define-random-distribution (biomial p n)
  "binomial distribution with parameters ~a and ~a"
  random-binomial
  binomial-pdf
  #f
  binomial-plot)

(define-random-distribution (geometric p)
  "geometric distribution with probability ~a"
  random-geometric
  geometric-pdf
  #f
  geometric-plot)

(define-random-distribution (logarithmic p)
  "logarithmic distribution with probability ~a"
  random-logarithmic
  logarithmic-pdf
  #f
  logarithmic-plot)

(define-random-distribution (poisson mu)
  "Poisson distribution with mean ~a"
  random-poisson
  poisson-pdf
  #f
  poisson-plot)

(define-random-distribution (discrete weights)
  "discrete distribution whose probability density is given by ~s"
  (lambda (s weights)
    (let ((d (make-discrete weights)))
      (random-discrete s d)))
  (lambda (x weights)
    (let ((d (make-discrete weights)))
      (discrete-pdf d x)))
  (lambda (x weights)
    (let ((d (make-discrete weights)))
      (discrete-cdf d x)))
  (lambda (weights)
    (let ((d (make-discrete weights)))
      (discrete-plot d))))

;;; Functions

;;; (random-distribution-description name) -> string?
;;;   name : symbol?
;;; Returns a general description of the random discription with the specified name.
(define (random-distribution-description name)
  (let ((dist (find-random-distribution name)))
    (apply format (random-distribution-form dist)
           (cdr (random-distribution-spec dist)))))

;;; (random-description spec) -> string?
;;;   spec : (cons/c symbol? (listof real?))
;;; Returns a description of the distribution specified by spec.
(define (random-description spec)
  (let ((dist (find-random-distribution (car spec))))
    (apply format (random-distribution-form dist) (cdr spec))))

;;; (random-sample spec [s]) -> real?
;;;   spec : (cons/c symbol? (listof any/c))
;;;   s : random-source? = (current-random-source)
;;; Returns a random sample from the distribution specified by spec using random
;;; source s. If s is not specified, (current-random-source) is used.
(define (random-sample spec (s (current-random-source)))
  (let ((dist (find-random-distribution (car spec))))
    (apply (random-distribution-sample dist) s (cdr spec))))

;;; (random-pdf-exists? spec) -> boolean?
;;;   spec : (cons/c symbol? (listof any/c))
;;; Returns true if a pdf function exists for the distribution specified by spec.
(define (random-pdf-exists? spec)
  (let ((dist (find-random-distribution (car spec))))
    (if (random-distribution-pdf dist) #t #f)))

;;; (random-pdf x spec) -> real?
;;;   x : real?
;;;   spec : (cons/c symbol? (listof any/c))
;;; Returns the value of the probability density function (pdf), p(x) for the
;;; distribution specified by spec.
(define (random-pdf x spec)
  (let ((dist (find-random-distribution (car spec))))
    (apply (random-distribution-pdf dist) x (cdr spec))))

;;; (random-cdf-exists? spec) -> boolean?
;;;   spec : (cons/c symbol? (listof any/c))
;;; Returns true if a cdf function exists for the distribution specified by spec.
(define (random-cdf-exists? spec)
  (let ((dist (find-random-distribution (car spec))))
    (if (random-distribution-cdf dist) #t #f)))

;;; (random-cdf x spec) -> (or/c real? false/c)
;;;   x : real?
;;;   spec : (cons/c symbol? (listof any/c))
;;; Returns the value of the cumulative density function (cdf), d(x) for the
;;; distribution specified by spec.
(define (random-cdf x spec)
  (let ((dist (find-random-distribution (car spec))))
    (apply (random-distribution-cdf dist) x (cdr spec))))

;;; (random-plot-exists? spec) -> boolean?
;;;   spec : (cons/c symbol? (listof any/c))
;;; Returns true if a plot function exists for the distribution specified by
;;; spec.
(define (random-plot-exists? spec)
  (let ((dist (find-random-distribution (car spec))))
    (if (random-distribution-pdf dist) #t #f)))

;;; (random-plot spec) -> any
;;;   spec : (cons/c symbol? (listof any/c))
;;; Returns a plot of the density function(s) for the distribution specified by
;;; spec.
(define (random-plot spec)
  (let ((dist (find-random-distribution (car spec))))
    (apply (random-distribution-plot dist) (cdr spec))))

;;; (random-sample-function spec) -> (->* () (random-source) () real?)
;;;   spec : (cons/c symbol? (listof any/c))
;;; Returns a function that returns a random sample from the distribution
;;; specified by spec. The returned function takes one optional argument, s, that
;;; specifies the random source. If s is not specified, (current-random-source)
;;; is used.
(define (random-sample-function spec)
  (let ((dist (find-random-distribution (car spec))))
    (lambda ((s (current-random-source)))
      (apply (random-distribution-sample dist) s (cdr spec)))))

;;; Module Contracts

(provide define-random-distribution)

(define random-spec/c
  (flat-named-contract
   'random-spec/c
   (lambda (x)
     (and (cons? x)
          (symbol? (car x))
          (list? (cdr x))
          (andmap real? (cdr x))))))
   

(provide/contract
 (random-distribution-description
  (-> symbol? string?))
 (random-description
  (-> random-spec/c string?))
 (random-sample
  (->* (random-spec/c) (random-source?) real?))
 (random-pdf-exists?
  (-> random-spec/c boolean?))
 (random-pdf
  (-> real? random-spec/c real?))
 (random-cdf-exists?
  (-> random-spec/c boolean?))
 (random-cdf
  (-> real? random-spec/c real?))
 (random-plot-exists?
  (-> random-spec/c boolean?))
 (random-plot
  (-> random-spec/c any))
 (random-sample-function
  (-> random-spec/c (->* () (random-source?) real?))))