;;;;;;;;;;
;;
;; Bayes network solver
;;
;; This file defines the data structures
;;
;;;;;;;;;;

;; A Bayes network (BNET) is a data structure that has the following components:
;; 1. a name, which is a symbol whose value is conventionally set to the BNET.
;; 2. a vector of the nodes in the network, each of which is a BN
;; 3. a dissection tree, which is the Cooper decomposition of the net

(define (make-bnet name nodes dissection-tree)
  (list name nodes dissection-tree))
(define bnet-name car)
(define bnet-nodes cadr)
(define bnet-dissection-tree caddr)
(define (set-bnet-name! x v) (set-car! x v))
(define (set-bnet-nodes! x v) (set-car! (cdr x) v))
(define (set-bnet-dissection-tree! x v) (set-car! (cddr x) v))

;; A node represents a random variable and its probabilistic relationships.
;; It is represented by a data structure that contains the following:
;; 1. a name
;; 2. list of parent nodes
;; 3. list of children nodes
;; 4. number of values that the node can take on
;; 5. vector of the actual values that the node can take on
;; 6. an array that holds the conditional probability distribution giving
;;    how this node depends on its parent nodes (see below for implementation)
;; 7. whether the node has a user-assigned specific value
;; 8. what that value is
;; 9. dissection records that must be reset when this node's value changes
;; 10. Markov blanket of this node
;; 11. neighborhood-size
;; We implement it with a binary tree.

(define f/t-vec (vector #f #t))

(define (make-bn name . rest)
  (let ((parents '())
        (children '())
        (n-values 2)
        (value-vector f/t-vec)
        (cdist '())
        (flagged #f)
        (pinned-value '())
        (indexed-records '())
        (Markov '())
        (neighborhood-size 0))
    (if rest
      (let ((rest1 (cdr rest)))
        (set! parents (car rest))
        (if rest1
          (let ((rest2 (cdr rest1)))
            (set! children (car rest1))
            (if rest2
              (let ((rest3 (cdr rest2)))
                (set! n-values (car rest2))
                (if rest3
                  (let ((rest4 (cdr rest3)))
                    (set! value-vector (car rest3))
                    (if rest4
                      (set! cdist (car rest4)))))))))))
    (cons (cons (cons name cdist)
                (cons parents children))
          (cons (cons (cons n-values value-vector)
                      (cons flagged pinned-value))
                (cons (cons Markov neighborhood-size)
                      indexed-records)))))

(define bn-name caaar)
(define bn-cdist cdaar)
(define bn-parents cadar)
(define bn-children cddar)
(define bn-n-values caaadr)
(define bn-value-vector cdaadr)
(define bn-flagged cadadr)
(define bn-pinned-value cddadr)
(define bn-Markov caaddr)
(define bn-neighborhood-size cdaddr)
(define bn-indexed-records cdddr)

(define (set-bn-name! n v) (set-car! (caar n) v))
(define (set-bn-cdist! n v) (set-cdr! (caar n) v))
(define (set-bn-parents! n v) (set-car! (cdar n) v))
(define (set-bn-children! n v) (set-cdr! (cdar n) v))
(define (set-bn-n-values! n v) (set-car! (caadr n) v))
(define (set-bn-value-vector! n v) (set-cdr! (caadr n) v))
(define (set-bn-flagged! n v) (set-car! (cdadr n) v))
(define (set-bn-pinned-value! n v) (set-cdr! (cdadr n) v))
(define (set-bn-Markov! n v) (set-car! (caddr n) v))
(define (set-bn-neighborhood-size! n v) (set-cdr! (caddr n) v))
(define (set-bn-indexed-records! n v) (set-cdr! (cddr n) v))

;; A dissection record is a recursive data structure that corresponds to an
;; efficient factoring of the network, according to Cooper's algorithm.

(define (make-rec summation-set evaluation-set instantiation-set
                  variable-set instantiation-cache net-y-ptr net-z-ptr)
  (cons (cons summation-set
              (cons evaluation-set instantiation-set))
        (cons (cons variable-set instantiation-cache)
              (cons net-y-ptr net-z-ptr))))
(define summation-set caar)
(define evaluation-set cadar)
(define instantiation-set cddar)
(define variable-set caadr)
(define instantiation-cache cdadr)
(define net-y-ptr caddr)
(define net-z-ptr cdddr)

;; Conditional probability distributions are represented by an n+1
;; dimensional array where a node has n parents.  The last dimension is the
;; index of the value of the node.  For example, a node X with parents Y and
;; Z, each of which is a binary variable with possible values #t or #f,
;; would have a three-dimensional table holding its conditional
;; probabilities.  Then, table[#t,#f,#f], or (table 1 0 0), would hold the
;; probability that X is #f given that Y is #t and Z is #f.  The order of
;; parents is the order in which they appear in bnet-nodes.
;; Because Scheme has no built-in support for n-dimensional arrays, we
;; implement them as n-nested vectors of vectors.