#lang racket/base

(provide
  distribution
  make-distribution
  distribution-generate
  distribution-add-to!
  distribution-contains?
  distribution-check-only
  distribution-linear-combination
  distribution-items
  distribution-filter
  distribution-empty?
  ; List format
  distribution->jsexpr
  jsexpr->distribution
  ; Object format
  distribution->jsexpr/object
  jsexpr->distribution/object
  )

; Distribution structure
; Distributions are list of items associated to a number indicating its frequency.
; When an element is picked from the list, the frequency of the items is taken into account.
(struct s-distribution
  ([total #:mutable] ; total number of elements in the distribution
   [lst-items #:mutable] ; same as items, but in a assoc-list
   items)) ; hash item -> number of occurences

; Exported function
(define distribution-items s-distribution-items)

; Add an item to the distribution, with its number of occurences
(define (distribution-add-to! d itm [additionnal-occurences 1])
  (define occurences (hash-ref (distribution-items d) itm 0))
  (hash-set! (s-distribution-items d) itm (+ occurences additionnal-occurences))
  (set-s-distribution-lst-items! d (hash->list (s-distribution-items d)))
  (set-s-distribution-total! d (+ (s-distribution-total d) additionnal-occurences)))

; Constructor with a list of item/occurences pairs
; item is (car pair), occurence is (cdr pair)
(define (distribution . pairs)
  (s-distribution
    (foldl
      (lambda (p res) (+ res (cdr p)))
      0
      pairs)
    pairs
    (make-hash pairs)))

; Syntax to build a distribution from the items and theirs occurences
; This syntax allow default values for items
; Note: the item is quoted
(define-syntax make-distribution
  (syntax-rules (*)
    ; No arguments - default constructor
    ((make-distribution)
     (s-distribution 0 (list) (make-hash)))
    ; default syntax
    ((make-distribution (id val) ...)
     (let ((dist (distribution)))
        (begin
          (distribution-add-to! dist (quote id) val) ...)
        dist))
    ; From a list and default values
    ((make-distribution lst (* default-val))
     (let ((dist (distribution)))
        (for-each
          (lambda (x) (distribution-add-to! dist x default-val))
          lst)
        dist))
    ; From a list, default values and values overriding the defaults
    ((make-distribution lst (* default-val) (id val) ...)
     (let ((h (make-hash))
           (dist (distribution)))
        (for-each
          (lambda (x) (hash-set! h x default-val))
          lst)
        (begin
          (hash-set! h (quote id) val) ...)
        (hash-for-each h (lambda (k v) (distribution-add-to! dist k v)))
        dist))
    ; With a useless list
    ((make-distribution lst (id val) ...)
     (let ((dist (distribution)))
        (begin
          (distribution-add-to! dist (quote id) val) ...)
        dist))
    ))

; Generate an item from a distribution
(define (distribution-generate d)
  ; Search function
  (define (search lst subtotal roll)
    (define new-subtotal (+ subtotal (cdar lst)))
    (if (< roll new-subtotal)
        (caar lst)
        (search (cdr lst) new-subtotal roll)))
  (search (s-distribution-lst-items d) 0 (random (s-distribution-total d))))

; Check if an item is in the distribution
(define (distribution-contains? d itm)
  (hash-has-key? (s-distribution-items d) itm))

; Check if there are unkown items in the distribution
; Raise an error if an unkown item is present
(define (distribution-check-only d known-items)
  (hash-for-each
    (s-distribution-items d)
    (lambda (k v)
      (when (not (member k known-items))
            (error "Distribution contains an unknown item: " k)))))

; Linear combination of distributions
; Make a new distribution from a dotted list of distributions and their weight
(define (distribution-linear-combination dists)
  (define d (distribution))
  (for-each
    (lambda (x)
      (define itms
        (hash-map
          (s-distribution-items (car x))
          (lambda (k v)
            (cons k (/ v (s-distribution-total (car x)))))))
      (for-each
        (lambda (itm)
          (distribution-add-to! d (car itm) (truncate (* (cdr x) (cdr itm) 100)))) ; * 100 and truncate because (cdr itm) is reduced
        itms))
    dists)
  d)

; Make a new distribution by filtering an existing distribution and keeping only the values that check a predicate
(define (distribution-filter d pred)
  (apply
    distribution
    (filter
      (lambda (x) (pred (car x)))
      (s-distribution-lst-items d))))

; Return true if the distribution is empty: no element in it
(define (distribution-empty? d)
  (eq? 0 (s-distribution-total d)))

; Serialisation to a json-compatible structure
; It is possible to pass a translator function for converting values that are not compatible to json
(define (distribution->jsexpr dist #:translator [tr (lambda (x) x)])
  (hash-map
    (s-distribution-items dist)
    (lambda (k v)
      (make-immutable-hash
        `((value . ,(tr k))
          (count . ,v))))))

; Object version: the distribution is serialised as an object whose attribute names are the distribution values and associated to their occurences.
; The translation function applies to the distribution values, and must return a symbol
(define (distribution->jsexpr/object dist #:translator [tr (lambda (x) x)])
  (define ret (make-hash))
  (hash-for-each
    (s-distribution-items dist)
    (lambda (k v)
      (hash-set! (tr k) v)))
  ret)

; Deserialisation from a json-compatible structure
; It is possible to pass a translator function for converting values that are not compatible to json
(define (jsexpr->distribution js #:translator [tr (lambda (x) x)])
  (define dist (distribution))
  (for-each
    (lambda (x)
      (distribution-add-to! dist (tr (hash-ref x 'value)) (hash-ref x 'count)))
    js)
  dist)

; Object version: the distribution is serialised as an object whose attribute names are the distribution values and associated to their occurences.
; The translation function applies to the distribution values, and must take a symbol in parameter
(define (jsexpr->distribution/object js #:translator [tr (lambda (x) x)])
  (define dist (distribution))
  (hash-for-each
    js
    (lambda (k v)
      (distribution-add-to! dist (tr k) v)))
  dist)
