; SPDX-FileCopyrightText: 2020 Arvydas Silanskas
; SPDX-FileCopyrightText: 2020 Bradley Lucier
; SPDX-FileCopyrightText: 2024 Peter McGoron
; SPDX-License-Identifier: MIT

(define-syntax proj1
  ;; Return the first returned value after evaluating `expr`.
  (syntax-rules ()
    ((proj1 expr)
     (let-values (((returned second) expr))
       returned))))

;;
;; This parameter is defined by the srfi-27 egg.
;;
#;(define current-random-source (make-parameter default-random-source))

(define (with-random-source random-source thunk)
  (unless (random-source? random-source)
    (error "expected random source"))
  (parameterize ((current-random-source random-source))
                (thunk)))

;;
;; Carefully return consecutive substreams of the s'th
;; SRFI 27 stream of random numbers.  See Sections 1.2 and
;; 1.3 of "An object-oriented random-number package with many
;; long streams and substreams", by Pierre L'Ecuyer, Richard
;; Simard, E. Jack Chen, and W. David Kelton, Operations Research,
;; vol. 50 (2002), pages 1073-1075.
;; https://doi.org/10.1287/opre.50.6.1073.358
;;

(: make-random-source-generator (integer -> (-> *)))
(define (make-random-source-generator s)
  (if (not (and (exact? s)
                (integer? s)
                (not (negative? s))))
      (error "make-random-source-generator: Expect nonnegative exact integer argument: " s)
      (let ((substream 0))
        (lambda ()
          (let ((new-source (make-random-source))) ;; deterministic
            (random-source-pseudo-randomize! new-source s substream)
            (set! substream (+ substream 1))
            new-source)))))

;;
;; Primitive randoms
;;

(: make-random-integer-generator (integer integer -> (-> integer)))
(define (make-random-integer-generator low-bound up-bound)
  (proj1 (make-uniform-random-integers high: (- up-bound 1)
                                       low: low-bound)))

(: make-random-u1-generator (-> (-> integer)))
(define (make-random-u1-generator)
  (make-random-integer-generator 0 2))

(: make-random-u8-generator (-> (-> integer)))
(define (make-random-u8-generator)
  (make-random-integer-generator 0 256))

(: make-random-s8-generator (-> (-> integer)))
(define (make-random-s8-generator)
  (make-random-integer-generator -128 128))

(: make-random-u16-generator (-> (-> integer)))
(define (make-random-u16-generator)
  (make-random-integer-generator 0 65536))

(: make-random-s16-generator (-> (-> integer)))
(define (make-random-s16-generator)
  (make-random-integer-generator -32768 32768))

(: make-random-u32-generator (-> (-> integer)))
(define (make-random-u32-generator)
  (make-random-integer-generator 0 (expt 2 32)))

(: make-random-s32-generator (-> (-> integer)))
(define (make-random-s32-generator)
  (make-random-integer-generator (- (expt 2 31)) (expt 2 31)))

(: make-random-u64-generator (-> (-> integer)))
(define (make-random-u64-generator)
  (make-random-integer-generator 0 (expt 2 64)))

(: make-random-s64-generator (-> (-> integer)))
(define (make-random-s64-generator)
  (make-random-integer-generator (- (expt 2 63)) (expt 2 63)))

(: clamp-real-number (number number number --> number))
(define (clamp-real-number lower-bound upper-bound value)
  (cond ((not (real? lower-bound))
         (error "expected real number for lower bound"))
        ((not (real? upper-bound))
         (error "expected real number for upper bound"))
        ((not (<= lower-bound upper-bound))
         (error "lower bound must be <= upper bound"))
        ((< value lower-bound) lower-bound)
        ((> value upper-bound) upper-bound)
        (else value)))

(: make-random-real-generator
   (number number -> (-> float)))
(define (make-random-real-generator low-bound up-bound)
  (unless (and (real? low-bound)
               (finite? low-bound))
    (error "expected finite real number for lower bound"))
  (unless (and (real? up-bound)
               (finite? up-bound))
     (error "expected finite real number for upper bound"))
  (unless (< low-bound up-bound)
     (error "lower bound must be < upper bound"))
  (let ((rand-real-proc (random-source-make-reals (current-random-source))))
   (lambda ()
     (define t (rand-real-proc))
     ;; alternative way of doing lowbound + t * (up-bound - low-bound)
     ;; is susceptible to rounding errors and would require clamping to be safe
     ;; (which in turn requires 144 for adjacent float function)
     (+ (* t low-bound)
        (* (- 1.0 t) up-bound)))))

(: make-random-rectangular-generator
   (number number number number -> (-> cplxnum)))
(define (make-random-rectangular-generator
          real-lower-bound real-upper-bound
          imag-lower-bound imag-upper-bound)
  (let ((real-gen (make-random-real-generator real-lower-bound real-upper-bound))
        (imag-gen (make-random-real-generator imag-lower-bound imag-upper-bound)))
    (lambda ()
      (make-rectangular (real-gen) (imag-gen)))))

;;; TODO: How to type optional arguments that appear in the front in
;;; CHICKEN?
(define make-random-polar-generator
  (case-lambda
    ((magnitude-lower-bound magnitude-upper-bound)
     (make-random-polar-generator 0+0i magnitude-lower-bound magnitude-upper-bound 0 (* 2 PI)))
    ((origin magnitude-lower-bound magnitude-upper-bound)
     (make-random-polar-generator origin magnitude-lower-bound magnitude-upper-bound 0 (* 2 PI)))
    ((magnitude-lower-bound magnitude-upper-bound angle-lower-bound angle-upper-bound)
     (make-random-polar-generator 0+0i magnitude-lower-bound magnitude-upper-bound angle-lower-bound angle-upper-bound))
    ((origin magnitude-lower-bound magnitude-upper-bound angle-lower-bound angle-upper-bound)
     (unless (complex? origin)
       (error "origin should be complex number"))
     (unless (and (real? magnitude-lower-bound)
                  (real? magnitude-upper-bound)
                  (real? angle-lower-bound)
                  (real? angle-upper-bound))
       (error "magnitude and angle bounds should be real numbers"))
     (unless (and (<= 0 magnitude-lower-bound)
                  (<= 0 magnitude-upper-bound))
       (error "magnitude bounds should be positive"))
     (unless (< magnitude-lower-bound magnitude-upper-bound)
       (error "magnitude lower bound should be less than upper bound"))
     (when (= angle-lower-bound angle-upper-bound)
       (error "angle bounds shouldn't be equal"))
     (let* ((b (square magnitude-lower-bound))
            (m (- (square magnitude-upper-bound) b))
            (t-gen (make-random-real-generator 0. 1.))
            (phi-gen (make-random-real-generator angle-lower-bound angle-upper-bound)))
       (lambda ()
         (let* ((t (t-gen))
                (phi (phi-gen))
                (r (sqrt (+ (* m t) b))))
          (+ origin (make-polar r phi))))))))

(: make-random-boolean-generator (-> (-> boolean)))
(define (make-random-boolean-generator)
  (proj1 (make-random-bernoullis)))

(: make-random-char-generator (string -> (-> char)))
(define (make-random-char-generator str)
  (when (not (string? str))
    (error "expected string"))
  (unless (> (string-length str) 0)
    (error "given string is of length 0"))
  (let* ((int-gen (make-random-integer-generator 0 (string-length str))))
   (lambda ()
     (string-ref str (int-gen)))))

(: make-random-string-generator (integer string -> (-> string)))
(define (make-random-string-generator k str)
  (let ((char-gen (make-random-char-generator str))
        (int-gen (make-random-integer-generator 0 k)))
    (lambda ()
      (generator->string char-gen (int-gen)))))

;;
;; Non-uniform distributions
;;

(: PI float)
(define PI (* 4 (atan 1.0)))

(: make-bernoulli-generator (#!optional number * -> (-> fixnum)))
(define make-bernoulli-generator
  (case-lambda
    (() (make-bernoulli-generator 1/2 (current-random-real)))
    ((p) (make-bernoulli-generator p (current-random-real)))
    ((p source)
     (gmap (lambda (x) (if x 1 0))
           (proj1 (make-random-bernoullis p: p randoms: source))))))

(: make-binomial-generator (number number -> (-> number)))
(define (make-binomial-generator n p)
  (proj1
   (make-random-binomials t: n p: p randoms: (current-random-real))))

(: make-categorical-generator
   ((vector-of number) -> (-> integer)))
(define (make-categorical-generator weights-vec)
  (define weight-sum
    (vector-fold
      (lambda (sum p)
        (unless (and (number? p)
                     (> p 0))
          (error "parameter must be a vector of positive numbers"))
        (+ sum p))
      0
      weights-vec))
  (define length (vector-length weights-vec))
  (let ((real-gen (make-random-real-generator 0 weight-sum)))
   (lambda ()
     (define roll (real-gen))
     (let it ((sum 0)
              (i 0))
       (define newsum (+ sum (vector-ref weights-vec i)))
       (if (or (< roll newsum)
               ;; in case of rounding errors and no matches, return last element
               (= i (- length 1)))
           i
           (it newsum
               (+ i 1)))))))

(: make-normal-generator
   (#!optional number number -> (-> number)))
(define make-normal-generator
  (case-lambda
    (() (make-normal-generator #i0 #i1))
    ((mean) (make-normal-generator mean #i1))
    ((mean deviation)
     (proj1
      (make-random-normals mu: mean sigma: deviation
                           randoms: (current-random-real))))))

;;; NOTE: Chicken's SRFI-27 has an exponential generator, but for
;;; some reason it is limited to means between 0 and 1.
(: make-exponential-generator (number -> (-> number)))
(define (make-exponential-generator mean)
  (unless (and (real? mean)
               (finite? mean)
               (positive? mean))
    (error "expected mean to be finite positive real number"))
  (let ((rand-real-proc (random-source-make-reals (current-random-source))))
   (lambda ()
     (- (* mean (log (rand-real-proc)))))))

(: make-geometric-generator (number -> (-> number )))
(define (make-geometric-generator p)
  (proj1
   (make-random-geometrics p: p randoms: (current-random-real))))

(: make-poisson-generator (number -> (-> number)))
(define (make-poisson-generator L)
  (proj1
   (make-random-poissons mu: L random: (current-random-real))))

(: gsampling (#!rest (-> *) -> (-> *)))
(define (gsampling . generators-lst)
  (let ((gen-vec (list->vector generators-lst))
        (rand-int-proc (random-source-make-integers (current-random-source))))

    ;remove exhausted generator at index
    (define (remove-gen index)
      (define new-vec (make-vector (- (vector-length gen-vec) 1)))
      ;when removing anything but first, copy all elements before index
      (when (> index 0)
        (vector-copy! new-vec 0 gen-vec 0 index))
      ;when removing anything but last, copy all elements after index
      (when (< index (- (vector-length gen-vec) 1))
        (vector-copy! new-vec index gen-vec (+ 1 index)))
      (set! gen-vec new-vec))

    ;randomly pick generator. If it's exhausted remove it, and pick again
    ;returns value (or eof, if all generators are exhausted)
    (define (pick)
      (let* ((index (rand-int-proc (vector-length gen-vec)))
             (gen (vector-ref gen-vec index))
             (value (gen)))
        (if (eof-object? value)
            (begin
              (remove-gen index)
              (if (= (vector-length gen-vec) 0)
                  (eof-object)
                  (pick)))
            value)))

    (lambda ()
      (if (= 0 (vector-length gen-vec))
          (eof-object)
          (pick)))))