#lang racket

;;; dds/networks

;;; This module provides some quick definitions for and analysing
;;; network models.  A network is a set of variables which are updated
;;; according to their corresponding update functions.  The variables
;;; to be updated at each step are given by the mode.
;;;
;;; This model can generalise Boolean networks, TBANs, multivalued
;;; networks, etc.

(require "utils.rkt" "generic.rkt" "functions.rkt"
         graph racket/random racket/hash)

(provide
 ;; Structures
 (contract-out [struct tbf/state ([weights (hash/c variable? number?)]
                                  [threshold number?])]
               [struct dynamics ([network network?]
                                 [mode mode?])])
 ;; Functions
 (contract-out [update (-> network? state? (set/c variable? #:kind 'dont-care) state?)]
               [make-state (-> (listof (cons/c symbol? any/c)) state?)]
               [make-state-booleanize (-> (listof (cons/c symbol? (or/c 0 1))) state?)]
               [booleanize-state (-> state? state?)]
               [make-network-from-functions (-> (listof (cons/c symbol? update-function/c)) network?)]
               [update-function-form->update-function (-> update-function-form? update-function/c)]
               [network-form->network (-> network-form? network?)]
               [make-network-from-forms (-> (listof (cons/c symbol? update-function-form?))
                                            network?)]
               [list-syntactic-interactions (-> network-form? variable? (listof variable?))]
               [build-syntactic-interaction-graph (-> network-form? graph?)]
               [build-all-states (-> domain-mapping/c (listof state?))]
               [make-same-domains (-> (listof variable?) generic-set? domain-mapping/c)]
               [make-boolean-domains (-> (listof variable?) (hash/c variable? (list/c #f #t)))]
               [make-01-domains (-> (listof variable?) (hash/c variable? (list/c 0 1)))]
               [build-all-boolean-states (-> (listof variable?) (listof state?))]
               [build-all-01-states (-> (listof variable?) (listof state?))]
               [get-interaction-sign (-> network? domain-mapping/c variable? variable? (or/c '+ '- '0))]
               [build-signed-interaction-graph/form (-> network-form? domain-mapping/c graph?)]
               [build-boolean-signed-interaction-graph/form (-> network-form? graph?)]
               [build-signed-interaction-graph (-> network? domain-mapping/c graph?)]
               [build-boolean-signed-interaction-graph (-> network? graph?)]
               [make-asyn (-> (listof variable?) mode?)]
               [make-syn (-> (listof variable?) mode?)]
               [make-dynamics-from-func (-> network? (-> (listof variable?) mode?) dynamics?)]
               [make-asyn-dynamics (-> network? dynamics?)]
               [make-syn-dynamics (-> network? dynamics?)]
               [read-org-network-make-asyn (-> string? dynamics?)]
               [read-org-network-make-syn (-> string? dynamics?)]
               [dds-step-one (-> dynamics? state? (set/c state?))]
               [dds-step-one-annotated (-> dynamics? state? (set/c (cons/c modality? state?)))]
               [dds-step (-> dynamics? (set/c state? #:kind 'dont-care) (set/c state?))]
               [dds-build-state-graph (-> dynamics? (set/c state? #:kind 'dont-care) graph?)]
               [dds-build-n-step-state-graph (-> dynamics? (set/c state? #:kind 'dont-care) number? graph?)]
               [dds-build-state-graph-annotated (-> dynamics? (set/c state? #:kind 'dont-care) graph?)]
               [dds-build-n-step-state-graph-annotated (-> dynamics? (set/c state? #:kind 'dont-care) number? graph?)]
               [pretty-print-state (-> state? string?)]
               [pretty-print-boolean-state (-> state? string?)]
               [pretty-print-state-graph-with (-> graph? (-> state? string?) graph?)]
               [pretty-print-state-graph (-> graph? graph?)]
               [ppsg (-> graph? graph?)]
               [pretty-print-boolean-state-graph (-> graph? graph?)]
               [ppsgb (-> graph? graph?)]
               [build-full-boolean-state-graph (-> dynamics? graph?)]
               [build-full-boolean-state-graph-annotated (-> dynamics? graph?)]
               [build-full-01-state-graph (-> dynamics? graph?)]
               [build-full-01-state-graph-annotated (-> dynamics? graph?)]
               [tabulate-state (->* (procedure? domain-mapping/c) (#:headers boolean?)
                                    (listof (listof any/c)))]
               [tabulate-state* (->* ((non-empty-listof procedure?) domain-mapping/c) (#:headers boolean?)
                                     (listof (listof any/c)))]
               [tabulate-state/boolean (->* (procedure? (listof variable?)) (#:headers boolean?)
                                            (listof (listof any/c)))]
               [tabulate-state*/boolean (->* ((non-empty-listof procedure?) (listof variable?)) (#:headers boolean?)
                                             (listof (listof any/c)))]
               [tabulate-network (->* (network? domain-mapping/c) (#:headers boolean?)
                                      (listof (listof any/c)))]
               [tabulate-boolean-network (->* (network?) (#:headers boolean?)
                                              (listof (listof any/c)))]
               [table->network (->* ((listof (*list/c any/c any/c))) (#:headers boolean?) network?)]
               [random-function/state (domain-mapping/c generic-set? . -> . procedure?)]
               [random-boolean-function/state ((listof variable?) . -> . procedure?)]
               [random-network (domain-mapping/c . -> . network?)]
               [random-boolean-network ((listof variable?) . -> . network?)]
               [random-boolean-network/vars (number? . -> . network?)]
               [apply-tbf-to-state (-> tbf? state? (or/c 0 1))]
               [tbf/state-w (-> tbf/state? (hash/c variable? number?))]
               [tbf/state-θ (-> tbf/state? number?)]
               [make-tbf/state (-> (listof (cons/c variable? number?)) number? tbf/state?)]
               [make-sbf/state (-> (listof (cons/c variable? number?)) sbf/state?)]
               [apply-tbf/state (-> tbf/state? (hash/c variable? (or/c 0 1)) (or/c 0 1))]
               [lists->tbfs/state (->* ((listof (listof (or/c number? symbol?))))
                                       (#:headers boolean?)
                                       (listof tbf/state?))]
               [lists->sbfs/state (->* ((listof (listof (or/c number? symbol?))))
                                       (#:headers boolean?)
                                       (listof sbf/state?))]
               [read-org-tbfs/state (->* (string?) (#:headers boolean?) (listof tbf/state?))]
               [read-org-sbfs/state (->* (string?) (#:headers boolean?) (listof sbf/state?))]
               [print-org-tbfs/state (->* ((non-empty-listof tbf/state?)) (#:headers boolean?)
                                          (listof (listof (or/c number? symbol?))))]
               [print-org-sbfs/state (->* ((non-empty-listof tbf/state?)) (#:headers boolean?)
                                          (listof (listof (or/c number? symbol?))))]
               [tbf/state-tabulate* (->* ((non-empty-listof tbf/state?)) (#:headers boolean?)
                                         (listof (listof (or/c symbol? number?))))]
               [tbf/state-tabulate (->* (tbf/state?) (#:headers boolean?)
                                        (listof (listof (or/c symbol? number?))))]
               [group-truth-table-by-nai (-> (listof (listof (or/c 0 1))) (listof (listof (listof (or/c 0 1)))))]
               [make-tbn (-> (listof (cons/c variable? tbf/state?)) tbn?)]
               [tbn->network (-> tbn? network?)]
               [make-sbn (-> (listof (cons/c variable? tbf/state?)) sbn?)]
               [parse-org-tbn (->* ((listof any/c))
                                   (#:headers boolean? #:func-names boolean?)
                                   tbn?)]
               [read-org-tbn (->* (string?)
                                  (#:headers boolean? #:func-names boolean?)
                                  tbn?)]
               [read-org-sbn (->* (string?)
                                  (#:headers boolean? #:func-names boolean?)
                                  tbn?)]
               [build-tbn-state-graph (-> tbn? graph?)]
               [normalized-tbn? (-> tbn? boolean?)]
               [normalize-tbn (-> tbn? normalized-tbn?)]
               [compact-tbf (-> tbf/state? tbf/state?)]
               [compact-tbn (-> tbn? tbn?)]
               [print-org-tbn (->* (tbn?) (#:headers boolean? #:func-names boolean?)
                                   (listof (listof (or/c number? symbol?))))]
               [print-org-sbn (->* (sbn?) (#:headers boolean? #:func-names boolean?)
                                   (listof (listof (or/c number? symbol?))))]
               [tbn-interaction-graph (->* (tbn?) (#:zero-edges boolean?)
                                           graph?)]
               [pretty-print-tbn-interaction-graph (-> graph? graph?)]
               [sbn-interaction-graph (-> sbn? graph?)])
 ;; Predicates
 (contract-out [variable? (-> any/c boolean?)]
               [state? (-> any/c boolean?)]
               [update-function-form? (-> any/c boolean?)]
               [network-form? (-> any/c boolean?)]
               [modality? (-> any/c boolean?)]
               [mode? (-> any/c boolean?)]
               [sbf/state? (-> any/c boolean?)])
 ;; Contracts
 (contract-out [state/c contract?]
               [update-function/c contract?]
               [domain-mapping/c contract?]
               [tbn? contract?]
               [sbn? contract?]))

(module+ test
  (require rackunit))


;;; =================
;;; Basic definitions
;;; =================

(define variable? symbol?)

;;; A state of a network is a mapping from the variables of the
;;; network to their values.
(define state? variable-mapping?)
(define state/c (flat-named-contract 'state state?))

;;; An update function is a function computing a value from the given
;;; state.
(define update-function/c (-> state? any/c))

;;; A network is a mapping from its variables to its update functions.
;;;
;;; Note that the domains of the variables of the network are not part
;;; of the network definition.  This is because the variables of some
;;; networks may have infinite domains, which can be restricted in
;;; multiple different ways.
(define network? (hash/c variable? procedure?))

;;; Given a state s updates all the variables from xs.
(define (update network s xs)
  (for/fold ([new-s s])
            ([x xs])
    (let ([f (hash-ref network x)])
      (hash-set new-s x (f s)))))

(module+ test
  (test-case "basic definitions"
    (define f1 (λ (s) (let ([x1 (hash-ref s 'x1)]
                            [x2 (hash-ref s 'x2)])
                        (and x1 (not x2)))))
    (define f2 (λ (s) (let ([x2 (hash-ref s 'x2)])
                        (not x2))))
    (define bn (make-network-from-functions `((x1 . ,f1) (x2 . ,f2))))
    (define s1 (make-state '((x1 . #t) (x2 . #f))))
    (define new-s1 (update bn s1 '(x2 x1)))
    (define s2 (make-state '((x1 . #f) (x2 . #f))))
    (define new-s2 (update bn s2 '(x2)))

    (check-equal? s1 #hash((x1 . #t) (x2 . #f)))
    (check-equal? new-s1 #hash((x1 . #t) (x2 . #t)))
    (check-equal? s2 #hash((x1 . #f) (x2 . #f)))
    (check-equal? new-s2 #hash((x1 . #f) (x2 . #t)))))

;;; A version of make-immutable-hash restricted to creating network
;;; states (see contract).
(define (make-state mappings) (make-immutable-hash mappings))

;;; Makes a new Boolean states from a state with numerical values 0
;;; and 1.
(define (make-state-booleanize mappings)
  (make-state (for/list ([mp mappings])
                (match mp
                  [(cons var 0) (cons var #f)]
                  [(cons var 1) (cons var #t)]))))

(module+ test
  (test-case "make-state, make-state-booleanize, booleanize-state"
    (check-equal? (make-state-booleanize '((a . 0) (b . 1)))
                  (make-state '((a . #f) (b . #t))))
    (check-equal? (booleanize-state (make-state '((a . 0) (b . 1))))
                  (make-state '((a . #f) (b . #t))))))

;;; Booleanizes a given state: replaces 0 with #f and 1 with #t.
(define (booleanize-state s)
  (for/hash ([(x val) s]) (match val [0 (values x #f)] [1 (values x #t)])))

;;; A version of make-immutable-hash restricted to creating networks.
(define (make-network-from-functions funcs) (make-immutable-hash funcs))


;;; =================================
;;; Syntactic description of networks
;;; =================================

;;; An update function form is any form which can appear as a body of
;;; a function and which can be evaluated with eval.  For example,
;;; '(and x y (not z)) or '(+ 1 a (- b 10)).
(define update-function-form? any/c)

;;; A Boolean network form is a mapping from its variables to the
;;; forms of their update functions.
(define network-form? variable-mapping?)

;;; Build an update function from an update function form.
(define (update-function-form->update-function form)
  (λ (s) (eval-with s form)))

(module+ test
  (test-case "update-function-form->update-function"
    (define s (make-state '((x . #t) (y . #f))))
    (define f (update-function-form->update-function '(and x y)))
    (check-equal? (f s) #f)))

;;; Build a network from a network form.
(define (network-form->network bnf)
  (for/hash ([(x form) bnf])
    (values x (update-function-form->update-function form))))

(module+ test
  (test-case "network-form->network"
    (define bn (network-form->network
         (make-hash '((a . (and a b)) (b . (not b))))))
    (define s (make-state '((a . #t) (b . #t))))
    (check-equal? ((hash-ref bn 'a) s) #t)))

;;; Build a network from a list of pairs of forms of update functions.
(define (make-network-from-forms forms)
  (network-form->network (make-immutable-hash forms)))

(module+ test
  (test-case "make-network-from-forms"
    (define bn (make-network-from-forms '((a . (and a b))
                                          (b . (not b)))))
    (define s (make-state '((a . #t) (b . #t))))
    (check-equal? ((hash-ref bn 'a) s) #t)))


;;; ============================
;;; Inferring interaction graphs
;;; ============================

;;; In this section I provide inference of both unsigned and signed
;;; interaction graphs, but since the inference of signed interaction
;;; graphs is based on analysing the dynamics of the networks, it may
;;; be quite resource-consuming, especially since I allow any
;;; syntactic forms in the definitions of the functions.

;;; Lists the variables of the network form appearing in the update
;;; function form for x.
(define (list-syntactic-interactions nf x)
  (set-intersect
   (extract-symbols (hash-ref nf x))
   (hash-keys nf)))

(module+ test
  (test-case "list-syntactic-interactions"
    (define n #hash((a . (+ a b c))
                    (b . (- b c))))
    (check-true (set=? (list-syntactic-interactions n 'a) '(a b)))
    (check-true (set=? (list-syntactic-interactions n 'b) '(b)))))

;;; Builds the graph in which the vertices are the variables of a
;;; given network, and which contains an arrow from a to b whenever a
;;; appears in (list-interactions a).
;;;
;;; Note that, while this definition is an easy one to check
;;; structurally, this is *not* how interaction graphs are typically
;;; defined.  An interaction graph is usually defined based on the
;;; dynamics of the network: an arrow from a variable x to a variable
;;; y means that varying x and only x may have an influence on the
;;; value of y.  It is easy to imagine a situation in which the
;;; syntactic interaction graph does not in fact agree with this
;;; criterion, the simplest example being the network y = x AND (NOT
;;; x).
(define (build-syntactic-interaction-graph n)
  (transpose
   (unweighted-graph/adj
    (for/list ([(var _) n]) (cons var (list-syntactic-interactions n var))))))

(module+ test
  (test-case "build-syntactic-interaction-graph"
    (define n #hash((a . (+ a b c))
             (b . (- b c))))
    (define ig (build-syntactic-interaction-graph n))
    (check-true (has-vertex? ig 'a))
    (check-true (has-vertex? ig 'b))
    (check-false (has-vertex? ig 'c))
    (check-true (has-edge? ig 'a 'a))
    (check-true (has-edge? ig 'b 'a))
    (check-true (has-edge? ig 'b 'b))
    (check-false (has-edge? ig 'c 'b))
    (check-false (has-edge? ig 'c 'a))))

;;; A domain mapping is a hash set mapping variables to the lists of
;;; values in their domains.
(define domain-mapping/c (hash/c variable? list?))

;;; Given a hash-set mapping variables to generic sets of their
;;; possible values, constructs the list of all possible states.
(define (build-all-states vars-domains)
  (let* ([var-dom-list (hash-map vars-domains (λ (x y) (cons x y)) #t)]
         [vars (map car var-dom-list)]
         [domains (map cdr var-dom-list)])
    (for/list ([s (apply cartesian-product domains)])
      (make-state (for/list ([var vars] [val s])
                    (cons var val))))))

(module+ test
  (test-case "build-all-states"
    (check-equal? (build-all-states #hash((a . (#t #f)) (b . (1 2 3))))
                  '(#hash((a . #t) (b . 1))
                    #hash((a . #t) (b . 2))
                    #hash((a . #t) (b . 3))
                    #hash((a . #f) (b . 1))
                    #hash((a . #f) (b . 2))
                    #hash((a . #f) (b . 3))))))

;;; Makes a hash set mapping all variables to a single domain.
(define (make-same-domains vars domain)
  (for/hash ([var vars]) (values var domain)))

;;; Makes a hash set mapping all variables to the Boolean domain.
(define (make-boolean-domains vars)
  (make-same-domains vars '(#f #t)))

(module+ test
  (test-case "make-same-domains, make-boolean-domains"
    (check-equal? (make-boolean-domains '(a b))
                  #hash((a . (#f #t)) (b . (#f #t))))))

;;; Makes a hash set mapping all variables to the Boolean domain,
;;; expressed as {0,1}.
(define (make-01-domains vars)
  (make-same-domains vars '(0 1)))

(module+ test
  (test-case "make-01-domains"
    (check-equal? (make-01-domains '(a b))
                  '#hash((a . (0 1)) (b . (0 1))))))

;;; Builds all boolean states possible over a given set of variables.
(define (build-all-boolean-states vars)
  (build-all-states (make-boolean-domains vars)))

(module+ test
  (test-case "build-all-boolean-states"
    (check-equal? (build-all-boolean-states '(a b))
                  '(#hash((a . #f) (b . #f))
                    #hash((a . #f) (b . #t))
                    #hash((a . #t) (b . #f))
                    #hash((a . #t) (b . #t))))))

;;; Builds all Boolean states over a given set of variables, but with
;;; 0 and 1 for Boolean values.
(define build-all-01-states
  (compose build-all-states make-01-domains))

(module+ test
  (test-case "build-all-01-states"
    (check-equal? (build-all-01-states '(a b))
                  '(#hash((a . 0) (b . 0))
                    #hash((a . 0) (b . 1))
                    #hash((a . 1) (b . 0))
                    #hash((a . 1) (b . 1))))))

;;; Given two interacting variables of a network and the domains
;;; of the variables, returns '+ if the interaction is monotonously
;;; increasing, '- if it is monotonously decreasing, and '0 otherwise.
;;;
;;; This function does not check whether the two variables indeed
;;; interact.  Its behaviour is undefined if the variables do not
;;; interact.
;;;
;;; /!\ This function iterates through almost all of the states of the
;;; network, so its performance decreases very quickly with network
;;; size.
(define (get-interaction-sign network doms x y)
  (let* ([dom-x (hash-ref doms x)]
         [dom-y (hash-ref doms y)]
         ;; Replace the domain of x by a dummy singleton.
         [doms-no-x (hash-set doms x '(#f))]
         ;; Build all the states, but as if x were not there: since I
         ;; replace its domain by a singleton, all states will contain
         ;; the same value for x.
         [states-no-x (build-all-states doms-no-x)]
         ;; Go through all states, then through all ordered pairs of
         ;; values of x, generate pairs of states (s1, s2) such that x
         ;; has a smaller value in s1, and check that updating y in s1
         ;; yields a smaller value than updating y in s2.  I rely on
         ;; the fact that the domains are ordered.
         [x-y-interactions (for*/list ([s states-no-x]
                                       [x1 dom-x] ; ordered pairs of values of x
                                       [x2 (cdr (member x1 dom-x))])
                             (let* ([s1 (hash-set s x x1)] ; s1(x) < s2(x)
                                    [s2 (hash-set s x x2)]
                                    [y1 ((hash-ref network y) s1)]
                                    [y2 ((hash-ref network y) s2)])
                               ;; y1 <= y2?
                               (<= (index-of dom-y y1) (index-of dom-y y2))))])
    (cond
      ;; If, in all interactions, y1 <= y2, then we have an
      ;; increasing/promoting interaction between x and y.
      [(andmap (λ (x) (eq? x #t)) x-y-interactions) '+]
      ;; If, in all interactions, y1 > y2, then we have an
      ;; decreasing/inhibiting interaction between x and y.
      [(andmap (λ (x) (eq? x #f)) x-y-interactions) '-]
      ;; Otherwise the interaction is neither increasing nor
      ;; decreasing.
      [else '0])))

(module+ test
  (test-case "get-interaction-sign"
    (define n #hash((a . (not b)) (b . a)))
    (define doms (make-boolean-domains '(a b)))
    (check-equal? (get-interaction-sign (network-form->network n) doms 'a 'b) '+)
    (check-equal? (get-interaction-sign (network-form->network n) doms 'b 'a) '-)))

;;; Constructs a signed interaction graph of a given network form,
;;; given the ordered domains of its variables.  The order on the
;;; domains determines the signs which will appear on the interaction
;;; graph.
;;;
;;; /!\ This function iterates through almost all states of the
;;; network for every arrow in the unsigned interaction graph, so its
;;; performance decreases very quickly with the size of the network.
(define (build-signed-interaction-graph/form network-form doms)
  (let ([ig (build-syntactic-interaction-graph network-form)]
        [network (network-form->network network-form)])
    ;; Label every edge of the interaction graph with the sign.
    (define sig
      (weighted-graph/directed
       (for/list ([e (in-edges ig)])
         (match-let ([(list x y) e])
           (list (get-interaction-sign network doms x y)
                 x y)))))
    ;; Ensure that every variable of the network appears in the signed
    ;; interaction graph as well.
    (for ([v (in-vertices ig)])
      (add-vertex! sig v))
    sig))

(module+ test
  (test-case "build-signed-interaction-graph/form"
    (define n #hash((a . (not b)) (b . a)))
    (define doms (make-boolean-domains '(a b)))
    (define sig1 (build-signed-interaction-graph/form n doms))
    (check-true (has-vertex? sig1 'a))
    (check-true (has-vertex? sig1 'b))
    (check-false (has-vertex? sig1 'c))
    (check-false (has-edge? sig1 'a 'a))
    (check-true (has-edge? sig1 'b 'a))
    (check-false (has-edge? sig1 'b 'b))
    (check-false (has-edge? sig1 'c 'b))
    (check-false (has-edge? sig1 'c 'a))
    (check-equal? (edge-weight sig1 'a 'b) '+)
    (check-equal? (edge-weight sig1 'b 'a) '-)))

;;; Calls build-signed-interaction-graph with the Boolean domain for
;;; all variable.
;;;
;;; /!\ The same performance warning applies as for
;;; build-signed-interaction-graph.
(define (build-boolean-signed-interaction-graph/form network-form)
  (build-signed-interaction-graph/form
   network-form
   (make-boolean-domains (hash-keys network-form))))

(module+ test
  (test-case "build-boolean-signed-interaction-graph/form"
    (define n #hash((a . (not b)) (b . a)))
    (define sig2 (build-boolean-signed-interaction-graph/form n))
    (check-true (has-vertex? sig2 'a))
    (check-true (has-vertex? sig2 'b))
    (check-false (has-vertex? sig2 'c))
    (check-false (has-edge? sig2 'a 'a))
    (check-true (has-edge? sig2 'b 'a))
    (check-false (has-edge? sig2 'b 'b))
    (check-false (has-edge? sig2 'c 'b))
    (check-false (has-edge? sig2 'c 'a))
    (check-equal? (edge-weight sig2 'a 'b) '+)
    (check-equal? (edge-weight sig2 'b 'a) '-)))

;;; Similar to build-signed-interaction-graph/form, but operates on a
;;; network rather than a form.  The resulting graph only includes the
;;; edges for positive or negative interactions.
;;;
;;; This function has operates with much less knowledge than
;;; build-signed-interaction-graph/form, so prefer using the latter
;;; when you can get a network form.
;;;
;;; /!\ This function iterates through all states of the network for
;;; every arrow in the unsigned interaction graph, so its performance
;;; decreases very quickly with the size of the network.
(define (build-signed-interaction-graph network doms)
  (define sig
    (weighted-graph/directed
     (for*/fold ([edges '()])
                ([(x _) (in-hash network)]
                 [(y _) (in-hash network)])
       (match (get-interaction-sign network doms x y)
         ['0 edges]
         [sign (cons (list sign x y) edges)]))))
  ;; Ensure that all variables of the network appear in the signed
  ;; interaction graph.
  (for ([(v _) (in-hash network)])
      (add-vertex! sig v))
    sig)

;;; Calls build-signed-interaction-graph assuming that the domains of
;;; all variables are Boolean.
;;;
;;; This function has operates with much less knowledge than
;;; build-boolean-signed-interaction-graph/form, so prefer using the
;;; latter when you can get a network form.
;;;
;;; /!\ This function iterates through all states of the network for
;;; every arrow in the unsigned interaction graph, so its performance
;;; decreases very quickly with the size of the network.
(define (build-boolean-signed-interaction-graph network)
  (build-signed-interaction-graph network (make-boolean-domains (hash-keys network))))

(module+ test
  (test-case "build-signed-interaction-graph, build-boolean-signed-interaction-graph"
    (define n #hash((a . (not b)) (b . a)))
    (define sig3 (build-boolean-signed-interaction-graph (network-form->network n)))
    (check-true (has-vertex? sig3 'a))
    (check-true (has-vertex? sig3 'b))
    (check-equal? (edge-weight sig3 'a 'a) '+)
    (check-equal? (edge-weight sig3 'b 'b) '+)
    (check-equal? (edge-weight sig3 'a 'b) '+)
    (check-equal? (edge-weight sig3 'b 'a) '-)))

;;; Interaction graphs for networks without interactions must still
;;; contain all nodes.
(module+ test
  (test-case "Interaction must graphs always contain all nodes."
    (define n #hash((a . #t) (b . #t)))
    (define ig (build-syntactic-interaction-graph n))
    (define sig-nf (build-boolean-signed-interaction-graph/form n))
    (define sig (build-boolean-signed-interaction-graph (network-form->network n)))
    (check-equal? (get-vertices ig) '(b a))
    (check-true (empty? (get-edges ig)))
    (check-equal? (get-vertices sig-nf) '(b a))
    (check-true (empty? (get-edges sig-nf)))
    (check-equal? (get-vertices sig) '(b a))))


;;; ====================
;;; Dynamics of networks
;;; ====================

;;; This section contains definitions for building and analysing the
;;; dynamics of networks.

;;; A modality is a set of variable.
(define modality? (set/c variable?))

;;; A mode is a set of modalities.
(define mode? (set/c modality?))

;;; A network dynamics is a network plus a mode.
(struct dynamics (network mode)
  #:methods gen:dds
  [;; Annotates each result state with the modality which lead to it.
   (define/match (dds-step-one-annotated dyn s)
     [((dynamics network mode) s)
      (for/set ([m mode]) (cons m (update network s m)))])])

;;; Given a list of variables, builds the asynchronous mode (a set of
;;; singletons).
(define (make-asyn vars)
  (for/set ([v vars]) (set v)))

;;; Given a list of variables, builds the synchronous mode (a set
;;; containing the set of variables).
(define (make-syn vars) (set (list->set vars)))

(module+ test
  (test-case "make-asyn, make-syn"
    (define vars '(a b c))
    (check-equal? (make-asyn vars) (set (set 'a) (set 'b) (set 'c)))
    (check-equal? (make-syn vars) (set (set 'a 'b 'c)))))

;;; Given a network, applies a function for building a mode to its
;;; variables and returns the corresponding network dynamics.
(define (make-dynamics-from-func network mode-func)
  (dynamics network (mode-func (hash-keys network))))

;;; Creates the asynchronous dynamics for a given network.
(define (make-asyn-dynamics network)
  (make-dynamics-from-func network make-asyn))

;;; Creates the synchronous dynamics for a given network.
(define (make-syn-dynamics network)
  (make-dynamics-from-func network make-syn))

(module+ test
  (test-case "make-asyn-dynamics, make-syn-dynamics"
    (define n (network-form->network #hash((a . (not a)) (b . b))))
    (define asyn (make-asyn-dynamics n))
    (define syn (make-syn-dynamics n))
    (check-equal? (dynamics-network asyn) n)
    (check-equal? (dynamics-mode asyn) (set (set 'a) (set 'b)))
    (check-equal? (dynamics-network syn) n)
    (check-equal? (dynamics-mode syn) (set (set 'a 'b)))))

;;; Reads an Org-mode-produced sexp, converts it into a network, and
;;; builds the asyncronous dynamics out of it.
(define read-org-network-make-asyn (compose make-asyn-dynamics network-form->network read-org-variable-mapping))

;;; Reads an Org-mode-produced sexp, converts it into a network, and
;;; builds the synchronous dynamics out of it.
(define read-org-network-make-syn (compose make-syn-dynamics network-form->network read-org-variable-mapping))

;;; Pretty-prints a state of the network.
(define (pretty-print-state s)
  (string-join (hash-map s (λ (key val) (format "~a:~a" key val)) #t)))

(module+ test
  (test-case "pretty-print-state"
    (check-equal? (pretty-print-state (make-state '((a . #f) (b . 3) (c . 4))))
                  "a:#f b:3 c:4")))

;;; Pretty-prints a state of the network to Boolean values 0 or 1.
(define (pretty-print-boolean-state s)
  (string-join (hash-map s (λ (key val) (format "~a:~a" key (any->01 val))) #t)))

(module+ test
  (test-case "pretty-print-boolean-state"
    (check-equal?
     (pretty-print-boolean-state (make-state '((a . #f) (b . #t) (c . #t))))
     "a:0 b:1 c:1")))

;;; Given a state graph and a pretty-printer for states build a new
;;; state graph with pretty-printed vertices and edges.
(define (pretty-print-state-graph-with gr pprinter)
  (update-graph gr #:v-func pprinter #:e-func pretty-print-set-sets))

;;; Pretty prints a state graph with pretty-print-state.
(define (pretty-print-state-graph gr)
  (pretty-print-state-graph-with gr pretty-print-state))

;;; A shortcut for pretty-print-state-graph.
(define ppsg pretty-print-state-graph)

;;; Pretty prints a state graph with pretty-print-boolean-state.
(define (pretty-print-boolean-state-graph gr)
  (pretty-print-state-graph-with gr pretty-print-boolean-state))

;;; A shortcut for pretty-print-boolean-state-graph.
(define ppsgb pretty-print-boolean-state-graph)

;;; Builds the full state graph of a Boolean network.
(define (build-full-boolean-state-graph dyn)
  (dds-build-state-graph
   dyn
   (list->set (build-all-boolean-states (hash-keys (dynamics-network dyn))))))

;;; Build the full annotated state graph of a Boolean network.
(define (build-full-boolean-state-graph-annotated dyn)
  (dds-build-state-graph-annotated
   dyn
   (list->set (build-all-boolean-states (hash-keys (dynamics-network dyn))))))

(module+ test
  (test-case "Dynamics of networks"
    (define n (network-form->network #hash((a . (not a)) (b . b))))
    (define asyn (make-asyn-dynamics n))
    (define syn (make-syn-dynamics n))
    (define s (make-state '((a . #t) (b . #f))))
    (define ss (set (make-state '((a . #t) (b . #t)))
                    (make-state '((a . #f) (b . #t)))))
    (define gr1 (dds-build-n-step-state-graph asyn (set s) 1))
    (define gr-full (dds-build-state-graph asyn (set s)))
    (define gr-full-pp (pretty-print-state-graph gr-full))
    (define gr-full-ppb (pretty-print-boolean-state-graph gr-full))
    (define gr-complete-bool (build-full-boolean-state-graph asyn))
    (define gr-complete-bool-ann (build-full-boolean-state-graph-annotated asyn))
    (check-equal? (dds-step-one asyn s) (set (make-state '((a . #f) (b . #f)))
                                             (make-state '((a . #t) (b . #f)))))
    (check-equal? (dds-step-one-annotated asyn s)
                  (set (cons (set 'b) '#hash((a . #t) (b . #f)))
                       (cons (set 'a) '#hash((a . #f) (b . #f)))))
    (check-equal? (dds-step-one syn s) (set (make-state '((a . #f) (b . #f)))))
    (check-equal? (dds-step asyn ss)
                  (set (make-state '((a . #f) (b . #t)))
                       (make-state '((a . #t) (b . #t)))))
    (check-true (has-vertex? gr1 #hash((a . #t) (b . #f))))
    (check-true (has-vertex? gr1 #hash((a . #f) (b . #f))))
    (check-false (has-vertex? gr1 #hash((a . #t) (b . #t))))
    (check-true (has-edge? gr1 #hash((a . #t) (b . #f)) #hash((a . #f) (b . #f))))
    (check-true (has-edge? gr1 #hash((a . #t) (b . #f)) #hash((a . #t) (b . #f))))
    (check-false (has-edge? gr1 #hash((a . #f) (b . #f)) #hash((a . #t) (b . #f))))

    (check-true (has-vertex? gr-full #hash((a . #t) (b . #f))))
    (check-true (has-vertex? gr-full #hash((a . #f) (b . #f))))
    (check-false (has-vertex? gr-full #hash((a . #t) (b . #t))))
    (check-true (has-edge? gr-full #hash((a . #t) (b . #f)) #hash((a . #f) (b . #f))))
    (check-true (has-edge? gr-full #hash((a . #t) (b . #f)) #hash((a . #t) (b . #f))))
    (check-true (has-edge? gr-full #hash((a . #f) (b . #f)) #hash((a . #t) (b . #f))))
    (check-true (has-edge? gr-full #hash((a . #f) (b . #f)) #hash((a . #f) (b . #f))))

    (check-true (has-vertex? gr-full-pp "a:#f b:#f"))
    (check-true (has-vertex? gr-full-pp "a:#t b:#f"))
    (check-true (has-vertex? gr-full-ppb "a:0 b:0"))
    (check-true (has-vertex? gr-full-ppb "a:1 b:0"))

    (check-true (set=?
                 (get-edges gr-complete-bool)
                 '((#hash((a . #f) (b . #f)) #hash((a . #t) (b . #f)))
                   (#hash((a . #f) (b . #f)) #hash((a . #f) (b . #f)))
                   (#hash((a . #t) (b . #f)) #hash((a . #t) (b . #f)))
                   (#hash((a . #t) (b . #f)) #hash((a . #f) (b . #f)))
                   (#hash((a . #t) (b . #t)) #hash((a . #f) (b . #t)))
                   (#hash((a . #t) (b . #t)) #hash((a . #t) (b . #t)))
                   (#hash((a . #f) (b . #t)) #hash((a . #f) (b . #t)))
                   (#hash((a . #f) (b . #t)) #hash((a . #t) (b . #t))))))

    (check-true (set=?
                 (get-edges gr-complete-bool-ann)
                 '((#hash((a . #f) (b . #f)) #hash((a . #t) (b . #f)))
                   (#hash((a . #f) (b . #f)) #hash((a . #f) (b . #f)))
                   (#hash((a . #t) (b . #f)) #hash((a . #t) (b . #f)))
                   (#hash((a . #t) (b . #f)) #hash((a . #f) (b . #f)))
                   (#hash((a . #t) (b . #t)) #hash((a . #f) (b . #t)))
                   (#hash((a . #t) (b . #t)) #hash((a . #t) (b . #t)))
                   (#hash((a . #f) (b . #t)) #hash((a . #f) (b . #t)))
                   (#hash((a . #f) (b . #t)) #hash((a . #t) (b . #t))))))
    (check-equal? (edge-weight gr-complete-bool-ann
                               #hash((a . #f) (b . #f)) #hash((a . #t) (b . #f)))
                  (set (set 'a)))
    (check-equal? (edge-weight gr-complete-bool-ann
                               #hash((a . #f) (b . #f)) #hash((a . #f) (b . #f)))
                  (set (set 'b)))
    (check-equal? (edge-weight gr-complete-bool-ann
                               #hash((a . #t) (b . #f)) #hash((a . #t) (b . #f)))
                  (set (set 'b)))
    (check-equal? (edge-weight gr-complete-bool-ann
                               #hash((a . #t) (b . #f)) #hash((a . #f) (b . #f)))
                  (set (set 'a)))
    (check-equal? (edge-weight gr-complete-bool-ann
                               #hash((a . #t) (b . #t)) #hash((a . #f) (b . #t)))
                  (set (set 'a)))
    (check-equal? (edge-weight gr-complete-bool-ann
                               #hash((a . #t) (b . #t)) #hash((a . #t) (b . #t)))
                  (set (set 'b)))
    (check-equal? (edge-weight gr-complete-bool-ann
                               #hash((a . #f) (b . #t)) #hash((a . #f) (b . #t)))
                  (set (set 'b)))
    (check-equal? (edge-weight gr-complete-bool-ann
                               #hash((a . #f) (b . #t)) #hash((a . #t) (b . #t)))
                  (set (set 'a)))))

;;; Like build-full-boolean-state-graph, but the states are expressed
;;; in 0 and 1, instead of #f and #t.
(define (build-full-01-state-graph dyn)
  (dds-build-state-graph
   dyn
   (list->set (build-all-01-states (hash-keys (dynamics-network dyn))))))

;;; Like build-full-boolean-state-graph-annotated, but the states are expressed
;;; in 0 and 1, instead of #f and #t.
(define (build-full-01-state-graph-annotated dyn)
  (dds-build-state-graph-annotated
   dyn
   (list->set (build-all-01-states (hash-keys (dynamics-network dyn))))))


;;; =================================
;;; Tabulating functions and networks
;;; =================================


;;; Like tabulate, but supposes that the function works on states.
;;;
;;; The argument domains defines the domains of each of the component
;;; of the states.  If headers it true, the resulting list starts with
;;; a listing the names of the variables of the domain and ending with
;;; the symbol 'f, which indicates the values of the function.
(define (tabulate-state func domains #:headers [headers #t])
  (define tab (tabulate-state* `(,func) domains #:headers headers))
  (cond
    [headers
     ;; Replace 'f1 in the headers by 'f.
     (match tab [(cons hdrs vals)
                 (cons (append (drop-right hdrs 1) '(f)) vals)])]
    [else tab]))

;;; Like tabulate-state, but assumes the function is a Boolean
;;; function.  args is a list of names of the arguments which can
;;; appear in the states.
(define (tabulate-state/boolean func args #:headers [headers #t])
  (tabulate-state func (make-boolean-domains args) #:headers headers))

(module+ test
  (test-case "tabulate-state/boolean"
    (define func (λ (st) (not (hash-ref st 'a))))
    (check-equal? (tabulate-state/boolean func '(a)) '((a f) (#f #t) (#t #f)))))

;;; Like tabulate-state, but takes a list of functions over the same
;;; domain.  If headers is #t, the first list of the result enumerates
;;; the variable names, and then contains a symbol 'fi for each of the
;;; functions, where i is replaced by the number of the function in
;;; the list.
(define (tabulate-state* funcs domains #:headers [headers #t])
  (define tab (for/list ([st (build-all-states domains)])
                (append (hash-map st (λ (x y) y) #t)
                        (for/list ([f funcs]) (f st)))))
  (cond
    [headers
     (define var-names (hash-map domains (λ (x y) x) #t))
     (define func-names (for/list ([_ funcs] [i (in-naturals 1)]) (string->symbol (format "f~a" i))))
     (cons (append var-names func-names) tab)]
    [else tab]))

;;; Like tabulate-state/boolean, but takes a list of functions.
(define (tabulate-state*/boolean funcs args #:headers [headers #t])
  (tabulate-state* funcs (make-boolean-domains args) #:headers headers))

(module+ test
  (test-case "tabulate-state*/boolean"
    (define f1 (λ (st) (and (hash-ref st 'a) (hash-ref st 'b))))
    (define f2 (λ (st) (or (hash-ref st 'a) (hash-ref st 'b))))
    (check-equal? (tabulate-state*/boolean (list f1 f2) '(a b))
                  '((a b f1 f2)
                    (#f #f #f #f)
                    (#f #t #f #t)
                    (#t #f #f #t)
                    (#t #t #t #t)))))

;;; Tabulates a given network.
;;;
;;; For a Boolean network with n variables, returns a table with 2n
;;; columns and 2^n rows.  The first n columns correspond to the
;;; different values of the variables of the networks.  The last n
;;; columns represent the values of the n update functions of the
;;; network.  If headers is #t, prepends a list of variable names and
;;; update functions (f-x, where x is the name of the corresponding
;;; variable) to the result.
(define (tabulate-network network domains #:headers [headers #t])
  ;; I use hash-map with try-order? set to #t to ask the hash table to
  ;; sort the keys for me.
  (define-values (vars funcs) (for/lists (l1 l2)
                                         ([pair (hash-map network cons #t)])
                                (values (car pair) (cdr pair))))
  (define tab (tabulate-state* funcs domains #:headers headers))
  (cond
    [headers
     ;; Replace the names of the functions tabulate-state* gave us by
     ;; what we promise in the comment.
     (define fnames (for/list ([x (in-list vars)])
                      (string->symbol (format "f-~a" x))))
     (match tab [(cons hdrs vals)
                 (cons (append (take hdrs (length vars)) fnames) vals)])]
    [else tab]))

;;; Like tabulate-network, but assumes all the variables are Boolean.
(define (tabulate-boolean-network bn #:headers [headers #t])
  (tabulate-network bn (make-boolean-domains (hash-map bn (λ (x y) x) #t))
                    #:headers headers))

(module+ test
  (test-case "tabulate-boolean-network"
    (define bn (network-form->network #hash((a . (not a)) (b . b))))
    (check-equal? (tabulate-boolean-network bn)
                  '((a b f-a f-b) (#f #f #t #f) (#f #t #t #t) (#t #f #f #f) (#t #t #f #t)))
    (check-equal? (tabulate-boolean-network bn #:headers #f)
                  '((#f #f #t #f) (#f #t #t #t) (#t #f #f #f) (#t #t #f #t)))))


;;; ===================================
;;; Constructing functions and networks
;;; ===================================

;;; Given a table like the one produced by tabulate-network,
;;; constructs a Boolean network having this behaviour.  If headers is
;;; #t, considers that the first element of the list are the headers
;;; and reads the names of the variables from them.  Otherwise
;;; generates names for variables of the form xi, where 0 ≤ i < number
;;; of variables, and treats all rows in the table as defining the
;;; behaviour of the functions of the network.  The columns defining
;;; the functions are taken to be in the same order as the variables
;;; in the first half of the function.  The headers of the columns
;;; defining the functions are therefore discarded.
;;;
;;; This function relies on table->function, so the same caveats
;;; apply.
(define (table->network table #:headers [headers #t])
  (define n (/ (length (car table)) 2))
  ;; Get the variable names from the table or generate them, if
  ;; necessary.
  (define var-names (cond [headers (take (car table) n)]
                          [else (for ([i (in-range n)])
                                  (symbol->string (format "x~a" i)))]))
  ;; Drop the headers if they are present.
  (define tab (cond [headers (cdr table)]
                    [else table]))
  ;; Split the table into the inputs and the outputs of the functions.
  (define-values (ins outs) (multi-split-at tab n))
  ;; Transpose outs to have functions define by lines instead of by
  ;; columns.
  (define func-lines (lists-transpose outs))
  ;; Make states out of inputs.
  (define st-ins (for/list ([in ins]) (make-state (map cons var-names in))))
  ;; Construct the functions.
  (define funcs (for/list ([out func-lines])
                  (table->function (for/list ([in st-ins] [o out])
                                     (list in o)))))
  ;; Construct the network.
  (make-network-from-functions (map cons var-names funcs)))

(module+ test
  (test-case "table->network"
    (define n (table->network '((x1 x2 f1 f2)
                                (#f #f #f #f)
                                (#f #t #f #t)
                                (#t #f #t #f)
                                (#t #t #t #t))))
    (define f1 (hash-ref n 'x1))
    (define f2 (hash-ref n 'x2))

    (check-false (f1 (make-state '((x1 . #f) (x2 . #f)))))
    (check-false (f1 (make-state '((x1 . #f) (x2 . #t)))))
    (check-true (f1 (make-state '((x1 . #t) (x2 . #f)))))
    (check-true (f1 (make-state '((x1 . #t) (x2 . #t)))))

    (check-false (f2 (make-state '((x1 . #f) (x2 . #f)))))
    (check-true (f2 (make-state '((x1 . #f) (x2 . #t)))))
    (check-false (f2 (make-state '((x1 . #t) (x2 . #f)))))
    (check-true (f2 (make-state '((x1 . #t) (x2 . #t)))))))


;;; =============================
;;; Random functions and networks
;;; =============================

;;; Generates a random function accepting a state over the domains
;;; given by arg-domains and producing values in func-domain.
(define (random-function/state arg-domains func-domain)
  (table->function (for/list ([st (build-all-states arg-domains)])
                     (list st (random-ref func-domain)))))

;;; Like random-function/state, but the domains of the arguments and
;;; of the function are Boolean.  args is a list of names of the
;;; variables appearing in the state.
(define (random-boolean-function/state args)
  (random-function/state (make-boolean-domains args) '(#f #t)))

(module+ test
  (test-case "random-boolean-function/state"
    (random-seed 0)
    (define f (random-boolean-function/state '(x1 x2)))
    (check-equal? (tabulate-state/boolean f '(x1 x2))
                  '((x1 x2 f) (#f #f #f) (#f #t #f) (#t #f #t) (#t #t #t)))
    (check-equal? (tabulate-state/boolean f '(x1 x2) #:headers #f)
                  '((#f #f #f) (#f #t #f) (#t #f #t) (#t #t #t)))
    (define bn (random-boolean-network/vars 3))
    (check-equal? (tabulate-boolean-network bn)
                  '((x0 x1 x2 f-x0 f-x1 f-x2)
                    (#f #f #f #f #t #f)
                    (#f #f #t #t #f #f)
                    (#f #t #f #f #t #t)
                    (#f #t #t #t #f #f)
                    (#t #f #f #t #f #t)
                    (#t #f #t #f #f #t)
                    (#t #t #f #f #f #f)
                    (#t #t #t #t #t #t)))))

;;; Generates a random network from the given domain mapping.
(define (random-network domains)
  (for/hash ([(x x-dom) (in-hash domains)])
    (values x (random-function/state domains x-dom))))

;;; Generates a random Boolean network with the given variables.
(define (random-boolean-network vars)
  (random-network (make-boolean-domains vars)))

;;; Like random-boolean-network, but also generates the names of the
;;; variables for the network.  The variables have the names x0 to xk,
;;; where k = n - 1.
(define (random-boolean-network/vars n)
  (random-boolean-network (for/list ([i (in-range n)]) (string->symbol (format "x~a" i)))))


;;; ===================
;;; TBF/TBN and SBF/SBN
;;; ===================

;;; Applies a TBF to a state.
;;;
;;; The values of the variables of the state are ordered by hash-map
;;; and fed to the TBF in order.  The number of the inputs of the TBF
;;; must match the number of variables in the state.
(define (apply-tbf-to-state tbf st)
  (apply-tbf tbf (list->vector (hash-map st (λ (_ val) val)))))

(module+ test
  (test-case "apply-tbf-to-state"
    (define st (make-state '((x1 . 0) (x2 . 1))))
    (define f (tbf #(1 1) 1))
    (check-equal? (apply-tbf-to-state f st) 0)))

;;; A state TBF is a TBF with named inputs.  A state TBF can be
;;; applied to states in an unambiguous ways.
(struct tbf/state (weights threshold) #:transparent)

;;; Shortcuts for acessing fields of a state/tbf.
(define tbf/state-w tbf/state-weights)
(define tbf/state-θ tbf/state-threshold)

;;; Makes a state/tbf from a list of pairs of names of variables and
;;; weights, as well as a threshold.
(define (make-tbf/state pairs threshold)
  (tbf/state (make-immutable-hash pairs) threshold))

(module+ test
  (test-case "tbf/state"
    (define f (make-tbf/state '((x1 . 1) (x2 . 1)) 1))
    (check-equal? (tbf/state-w f) #hash((x1 . 1) (x2 . 1)))
    (check-equal? (tbf/state-θ f) 1)))

;;; A sign Boolean function (SBF) is a TBF whose threshold is 0.
(define sbf/state? (and/c tbf/state? (λ (tbf) (zero? (tbf/state-θ tbf)))))

(module+ test
  (test-case "sbf/state?"
    (check-true (sbf/state? (tbf/state #hash((a . -1) (b . 1)) 0)))))

;;; Makes a state/tbf which is an SBF from a list of pairs of names of
;;; variables and weights.
(define (make-sbf/state pairs)
  (make-tbf/state pairs 0))

(module+ test
  (test-case "make-sbf/state"
    (check-equal? (make-sbf/state '((a . -1) (b . 1)))
                  (make-tbf/state '((a . -1) (b . 1)) 0))))

;;; Applies a state TBF to its inputs.
;;;
;;; Applying a TBF consists in multiplying the weights by the
;;; corresponding inputs and comparing the sum of the products to the
;;; threshold.
;;;
;;; This function is similar to apply-tbf, but applies a state TBF (a
;;; TBF with explicitly named inputs) to a state whose values are 0
;;; and 1.
(define (apply-tbf/state tbf st)
  (any->01 (> (foldl + 0 (hash-values
                          (hash-intersect (tbf/state-w tbf)
                                          st
                                          #:combine *)))
              (tbf/state-θ tbf))))

(module+ test
  (test-case "apply-tbf/state"
    (define st1 (make-state '((a . 1) (b . 0) (c . 1))))
    (define st2 (make-state '((a . 1) (b . 1) (c . 0))))
    (define tbf (make-tbf/state '((a . 2) (b . -2)) 1))
    (check-equal? (apply-tbf/state tbf st1) 1)
    (check-equal? (apply-tbf/state tbf st2) 0)))

;;; Reads a list of tbf/state from a list of list of numbers.
;;;
;;; The last element of each list is taken to be the threshold of the
;;; TBFs, and the rest of the elements are taken to be the weights.
;;;
;;; If headers is #t, the names of the variables to appear as the
;;; inputs of the TBF are taken from the first list.  The last element
;;; of this list is discarded.
;;;
;;; If headers is #f, the names of the variables are generated as xi,
;;; where i is the index of the variable.
(define (lists->tbfs/state lsts #:headers [headers #t])
  (define-values (var-names rows)
    (if headers
        (values (car lsts) (cdr lsts))
        (values (for/list ([i (in-range (length (car lsts)))])
                  (string->symbol (format "x~a" i)))
                lsts)))
  (for/list ([lst (in-list rows)])
    (define-values (ws θ) (split-at-right lst 1))
    (make-tbf/state (for/list ([x (in-list var-names)]
                               [w (in-list ws)])
                      (cons x w))
                    (car θ))))

(module+ test
  (test-case "lists->tbfs/state"
    (define tbfs '((1 2 3) (1 1 2)))
    (check-equal? (lists->tbfs/state tbfs #:headers #f)
                  (list
                   (tbf/state '#hash((x0 . 1) (x1 . 2)) 3)
                   (tbf/state '#hash((x0 . 1) (x1 . 1)) 2)))
    (check-equal? (lists->tbfs/state (cons '(a b f) tbfs))
                  (list
                   (tbf/state '#hash((a . 1) (b . 2)) 3)
                   (tbf/state '#hash((a . 1) (b . 1)) 2)))))

;;; Like lists->tbfs/state, but does not expect thresholds in the
;;; input.
;;;
;;; Every lists in the list contains the weights of the SBF.  If
;;; headers is #t, the names of the variables to appear as the inputs
;;; of the TBF are taken from the first list.
;;;
;;; If headers is #f, the names of the variables are generated as xi,
;;; where i is the index of the variable.
(define (lists->sbfs/state lsts #:headers [headers #t])
  (define rows (if headers (cdr lsts) lsts))
  (define rows-θ (for/list ([lst (in-list rows)]) (append lst '(0))))
  (lists->tbfs/state (if headers (cons (car lsts) rows-θ) rows-θ)
                     #:headers headers))

(module+ test
  (test-case "lists->sbfs/state"
    (define tbfs '((1 2) (1 -1)))
    (check-equal? (lists->sbfs/state tbfs #:headers #f)
                  (list
                   (tbf/state '#hash((x0 . 1) (x1 . 2)) 0)
                   (tbf/state '#hash((x0 . 1) (x1 . -1)) 0)))
    (check-equal? (lists->sbfs/state (cons '(a b) tbfs) #:headers #t)
                  (list
                   (tbf/state '#hash((a . 1) (b . 2)) 0)
                   (tbf/state '#hash((a . 1) (b . -1)) 0)))))

;;; Reads a list of tbf/state from an Org-mode string containing a
;;; sexp, containing a list of lists of numbers.  As in
;;; lists->tbfs/state, the last element of each list is taken to be
;;; the threshold of the TBFs, and the rest of the elements are taken
;;; to be the weights.
;;;
;;; If headers is #t, the names of the variables to appear as the
;;; inputs of the TBF are taken from the first list.  The last element
;;; of this list is discarded.
;;;
;;; If headers is #f, the names of the variables are generated as xi,
;;; where i is the index of the variable.
(define (read-org-tbfs/state str #:headers [headers #t])
  (lists->tbfs/state (read-org-sexp str) #:headers headers))

(module+ test
  (test-case "read-org-tbfs/state"
    (check-equal? (read-org-tbfs/state "((a b f) (1 2 3) (1 1 2))")
                  (list
                   (tbf/state '#hash((a . 1) (b . 2)) 3)
                   (tbf/state '#hash((a . 1) (b . 1)) 2)))
    (check-equal? (read-org-tbfs/state "((1 2 3) (1 1 2))" #:headers #f)
                  (list
                   (tbf/state '#hash((x0 . 1) (x1 . 2)) 3)
                   (tbf/state '#hash((x0 . 1) (x1 . 1)) 2)))))

;;; Like read-org-tbfs/state, but reads a list of SBFs.  Therefore,
;;; the lists of numbers in the sexp are taken to be the weights of
;;; the SBFs.
;;;
;;; If headers is #t, the names of the variables to appear as the
;;; inputs of the TBF are taken from the first list.  If headers is
;;; #f, the names of the variables are generated as xi, where i is the
;;; index of the variable.
(define (read-org-sbfs/state str #:headers [headers #t])
  (lists->sbfs/state (read-org-sexp str) #:headers headers))

(module+ test
  (test-case "read-org-sbfs/state"
    (check-equal? (read-org-sbfs/state "((a b) (-1 2) (1 1))")
                  (list
                   (tbf/state '#hash((a . -1) (b . 2)) 0)
                   (tbf/state '#hash((a . 1) (b . 1)) 0)))
    (check-equal? (read-org-sbfs/state "((-1 2) (1 1))" #:headers #f)
                  (list
                   (tbf/state '#hash((x0 . -1) (x1 . 2)) 0)
                   (tbf/state '#hash((x0 . 1) (x1 . 1)) 0)))))

;;; Given a list of tbf/state, produces a sexp that Org-mode can
;;; interpret as a table.
;;;
;;; All tbf/state in the list must have the same inputs.  The function
;;; does not check this property.
;;;
;;; If #:headers is #f, does not print the names of the inputs of the
;;; TBFs.  If #:headers is #t, the output starts by a list giving the
;;; names of the variables, as well as the symbol 'θ to represent the
;;; column giving the thresholds of the TBF.
(define (print-org-tbfs/state tbfs #:headers [headers #t])
  (define table (for/list ([tbf (in-list tbfs)])
                  (append (hash-map (tbf/state-w tbf) (λ (_ w) w) #t)
                          (list (tbf/state-θ tbf)))))
  (if headers
      (cons (append (hash-map (tbf/state-w (car tbfs)) (λ (x _) x) #t) '(θ))
            table)
      table))

(module+ test
  (test-case "print-org-tbfs/state"
    (define tbfs (list (make-tbf/state '((a . 1) (b . 2)) 3)
                       (make-tbf/state '((a . -2) (b . 1)) 1)))
    (check-equal? (print-org-tbfs/state tbfs)
                  '((a b θ) (1 2 3) (-2 1 1)))))

;;; Like print-org-tbfs/state, but expects a list of SBFs.  The
;;; thresholds are therefore not included in the output.
;;;
;;; All sbf/state in the list must have the same inputs.  The function
;;; does not check this property.
;;;
;;; If #:headers is #f, does not print the names of the inputs of the
;;; TBFs.  If #:headers is #t, the output starts by a list giving the
;;; names of the variables.
(define (print-org-sbfs/state sbfs #:headers [headers #t])
  (define table (for/list ([sbf (in-list sbfs)])
                  (hash-map (tbf/state-w sbf) (λ (_ w) w) #t)))
  (if headers
      (cons (hash-map (tbf/state-w (car sbfs)) (λ (x _) x) #t)
            table)
      table))

(module+ test
  (define sbfs (list (make-sbf/state '((a . 1) (b . 2)))
                     (make-sbf/state '((a . -2) (b . 1)))))
  (check-equal? (print-org-sbfs/state sbfs)
                '((a b) (1 2) (-2 1)))
  (check-equal? (print-org-sbfs/state sbfs #:headers #f)
                '((1 2) (-2 1))))

;;; Tabulates a list of tbf/state.
;;;
;;; As in the case of tbf-tabulate*, the result is a list of lists
;;; giving the truth tables of the given TBFs.  The first elements of
;;; each row give the values of the inputs, while the last elements
;;; give the values of each function corresponding to the input.
;;;
;;; All the TBFs must have exactly the same inputs.  This function
;;; does not check this property.
;;;
;;; If #:headers is #t, the output starts by a list giving the names
;;; of the variables, and then the symbols 'fi, where i is the number
;;; of the TBF in the list.
(define (tbf/state-tabulate* tbfs #:headers [headers #t])
  (define vars (hash-map (tbf/state-w (car tbfs)) (λ (x _) x) #t))
  (tabulate-state* (map (curry apply-tbf/state) tbfs)
                   (make-same-domains vars '(0 1))
                   #:headers headers))

(module+ test
  (test-case "tbf/state-tabulate*"
    (define tbfs (list (make-tbf/state '((a . 1) (b . 2)) 1)
                       (make-tbf/state '((a . -2) (b . 3)) 1)))
    (check-equal? (tbf/state-tabulate* tbfs)
                  '((a b f1 f2)
                    (0 0 0 0)
                    (0 1 1 1)
                    (1 0 0 0)
                    (1 1 1 0)))))

;;; Like tbf/state-tabulate*, but only tabulates a single TBF.
(define (tbf/state-tabulate tbf #:headers [headers #t])
  (tbf/state-tabulate* (list tbf) #:headers headers))

(module+ test
  (test-case "tbf/state-tabulate"
    (define tbf (make-tbf/state '((a . -2) (b . 3)) 1))
    (check-equal? (tbf/state-tabulate tbf)
                  '((a b f1)
                    (0 0 0)
                    (0 1 1)
                    (1 0 0)
                    (1 1 0)))))

;;; Given a truth table of a Boolean function, groups the lines by the
;;; "number of activated inputs"—the number of inputs which are 1 in
;;; the input vector.
;;;
;;; The truth table must not include the header line.
(define (group-truth-table-by-nai tt)
  (define sum (((curry foldl) +) 0))
  (group-by (λ (row) (drop-right row 1))
            tt
            (λ (in1 in2) (= (sum in1) (sum in2)))))

(module+ test
  (test-case "group-truth-table-by-nai"
    (check-equal? (group-truth-table-by-nai '((0 0 0 1)
                                              (0 0 1 1)
                                              (0 1 0 0)
                                              (0 1 1 1)
                                              (1 0 0 0)
                                              (1 0 1 0)
                                              (1 1 0 1)
                                              (1 1 1 0)))
                  '(((0 0 0 1))
                    ((0 0 1 1) (0 1 0 0) (1 0 0 0))
                    ((0 1 1 1) (1 0 1 0) (1 1 0 1))
                    ((1 1 1 0))))))

;;; A TBN is a network form mapping variables to tbf/state.
;;;
;;; The tbf/state must only reference variables appearing in the
;;; network.  This contract does not check this condition.
(define tbn? (hash/c variable? tbf/state?))

;;; Builds a TBN from a list of pairs (variable, tbf/state).
(define make-tbn make-immutable-hash)

(module+ test
  (test-case "make-tbn"
    (define tbf-not (make-tbf/state '((a . -1)) -1))
    (define tbf-id (make-sbf/state '((a . 1))))
    (check-equal? (make-tbn `((a . ,tbf-not) (b . ,tbf-id)))
                  (hash 'a (tbf/state '#hash((a . -1)) -1)
                        'b (tbf/state '#hash((a . 1)) 0)))))

;;; A SBN is a network form mapping variables to sbf/state.
;;;
;;; The tbf/state must only reference variables appearing in the
;;; network.  This contract does not check this condition.
(define sbn? (hash/c variable? sbf/state?))

;;; Builds an SBN from a list of pairs (variable, sbf/state).
(define make-sbn make-immutable-hash)

(module+ test
  (test-case "make-sbn"
    (define sbf1 (make-sbf/state '((a . -1))))
    (define sbf2 (make-sbf/state '((a . 1))))
    (check-equal? (make-sbn `((a . ,sbf1) (b . ,sbf2)))
                  (hash 'a (tbf/state '#hash((a . -1)) 0)
                        'b (tbf/state '#hash((a . 1)) 0)))))

;;; Constructs a network from a network form defining a TBN.
(define (tbn->network tbn)
  (for/hash ([(var tbf) (in-hash tbn)])
    (values var ((curry apply-tbf/state) tbf))))

(module+ test
  (test-case "tbn->network"
    (define tbn (make-tbn `((a . ,(make-sbf/state '((b . 1))))
                            (b . ,(make-tbf/state '((a . -1)) -1)))))
    (define n (tbn->network tbn))
    (define s1 (make-state '((a . 0) (b . 0))))
    (check-equal? (update n s1 '(a b))
                  (make-state '((a . 0) (b . 1))))

    (define sbn (make-sbn `((a . ,(make-sbf/state '((b . -1))))
                            (b . ,(make-sbf/state '((a . 1)))))))
    (define sn (tbn->network sbn))
    (define s2 (make-state '((a . 1) (b . 1))))
    (check-equal? (update sn s2 '(a b))
                  (make-state '((a . 0) (b . 1))))))

;;; A helper function for read-org-tbn and read-org-sbn.  It reads a
;;; TBN from an Org-mode sexp containing a list of lists of numbers.
;;; As in lists->tbfs/state, the last element of each list is taken to
;;; be the threshold of the TBFs, and the rest of the elements are
;;; taken to be the weights.
;;;
;;; As in read-org-tbfs/state, if headers is #t, the names of the
;;; variables to appear as the inputs of the TBF are taken from the
;;; first list.  The last element of this list is discarded.
;;; If headers is #f, the names of the variables are generated as xi,
;;; where i is the index of the variable.
;;;
;;; If func-names is #t, the first element in every row except the
;;; first one, are taken to be the name of the variable to which the
;;; TBF should be associated.  If func-names is #f, the functions are
;;; assigned to variables in alphabetical order.
;;;
;;; func-names cannot be #t if headers is #f.  The function does not
;;; check this condition.
(define (parse-org-tbn sexp
                       #:headers [headers #t]
                       #:func-names [func-names #t])
  (cond
    [(eq? func-names #t)
     (define-values (vars rows) (multi-split-at sexp 1))
     (define tbfs (lists->tbfs/state rows #:headers headers))
     (for/hash ([tbf (in-list tbfs)] [var (in-list (cdr vars))])
       (values (car var) tbf))]
    [else
     (define tbfs (lists->tbfs/state sexp #:headers headers))
     (define vars (hash-map (tbf/state-w (car tbfs)) (λ (x _) x) #t))
     (for/hash ([tbf (in-list tbfs)] [var (in-list vars)])
       (values var tbf))]))

;;; Reads a TBN from an Org-mode string containing a sexp, containing
;;; a list of lists of numbers.  As in lists->tbfs/state, the last
;;; element of each list is taken to be the threshold of the TBFs, and
;;; the rest of the elements are taken to be the weights.
;;;
;;; As in read-org-tbfs/state, if headers is #t, the names of the
;;; variables to appear as the inputs of the TBF are taken from the
;;; first list.  The last element of this list is discarded.
;;; If headers is #f, the names of the variables are generated as xi,
;;; where i is the index of the variable.
;;;
;;; If func-names is #t, the first element in every row except the
;;; first one, are taken to be the name of the variable to which the
;;; TBF should be associated.  If func-names is #f, the functions are
;;; assigned to variables in alphabetical order.
;;;
;;; func-names cannot be #t if headers is #f.  The function does not
;;; check this condition.
(define (read-org-tbn str
                      #:headers [headers #t]
                      #:func-names [func-names #t])
  (parse-org-tbn (read-org-sexp str)
                 #:headers headers
                 #:func-names func-names))

(module+ test
  (test-case "read-org-tbn, parse-org-tbn"
    (check-equal? (read-org-tbn "((\"-\" \"x\" \"y\" \"θ\") (\"y\" -1 0 -1) (\"x\" 0 -1 -1))")
                  (hash
                   'x
                   (tbf/state '#hash((x . 0) (y . -1)) -1)
                   'y
                   (tbf/state '#hash((x . -1) (y . 0)) -1)))
    (check-equal? (read-org-tbn "((\"x\" \"y\" \"θ\") (-1 0 -1) (0 -1 -1))" #:func-names #f)
                  (hash
                   'x
                   (tbf/state '#hash((x . -1) (y . 0)) -1)
                   'y
                   (tbf/state '#hash((x . 0) (y . -1)) -1)))
    (check-equal? (read-org-tbn "((-1 0 -1) (0 -1 -1))" #:headers #f #:func-names #f)
                  (hash
                   'x0
                   (tbf/state '#hash((x0 . -1) (x1 . 0)) -1)
                   'x1
                   (tbf/state '#hash((x0 . 0) (x1 . -1)) -1)))))

;;; Like read-org-tbn, but reads an SBN from an Org-mode string
;;; containing a sexp, containing a list of lists of numbers.
;;;
;;; As in read-org-sbfs/state, if headers is #t, the names of the
;;; variables to appear as the inputs of the SBF are taken from the
;;; first list.  The last element of this list is discarded.
;;; If headers is #f, the names of the variables are generated as xi,
;;; where i is the index of the variable.
;;;
;;; If func-names is #t, the first element in every row except the
;;; first one, are taken to be the name of the variable to which the
;;; TBF should be associated.  If func-names is #f, the functions are
;;; assigned to variables in alphabetical order.
;;;
;;; func-names cannot be #t if headers is #f.  The function does not
;;; check this condition.
(define (read-org-sbn str
                      #:headers [headers #t]
                      #:func-names [func-names #t])
  (define sexp (read-org-sexp str))
  ;; Inject the 0 thresholds into the rows of the sexp we have just read.
  (define (inject-0 rows) (for/list ([row (in-list rows)]) (append row '(0))))
  (define sexp-ready (if headers
                         (cons (car sexp) (inject-0 (cdr sexp)))
                         (inject-0 sexp)))
  (parse-org-tbn sexp-ready #:headers headers #:func-names func-names))

(module+ test
  (test-case "read-org-sbn, parse-org-tbn"
    (check-equal? (read-org-sbn "((\"-\" \"x\" \"y\") (\"y\" -1 0) (\"x\" 0 -1))")
                  (hash
                   'x
                   (tbf/state '#hash((x . 0) (y . -1)) 0)
                   'y
                   (tbf/state '#hash((x . -1) (y . 0)) 0)))
    (check-equal? (read-org-sbn "((\"x\" \"y\") (-1 0) (0 -1))" #:func-names #f)
                  (hash
                   'x
                   (tbf/state '#hash((x . -1) (y . 0)) 0)
                   'y
                   (tbf/state '#hash((x . 0) (y . -1)) 0)))
    (check-equal? (read-org-sbn "((-1 0) (0 -1))" #:headers #f #:func-names #f)
                  (hash
                   'x0
                   (tbf/state '#hash((x0 . -1) (x1 . 0)) 0)
                   'x1
                   (tbf/state '#hash((x0 . 0) (x1 . -1)) 0)))))

;;; A shortcut for building the state graphs of TBN.
(define build-tbn-state-graph
  (compose pretty-print-state-graph
           build-full-01-state-graph
           make-syn-dynamics
           tbn->network))

;;; Checks whether a TBN is normalized: whether all of the functions
;;; have the same inputs, and whether these inputs are exactly the
;;; variables of the TBN.
(define (normalized-tbn? tbn)
  (define tbn-vars (hash-keys tbn))
  (for/and ([tbf (in-list (hash-values tbn))])
    (set=? tbn-vars (hash-keys (tbf/state-w tbf)))))

(module+ test
  (test-case "normalized-tbn?"
    (check-false (normalized-tbn?
                  (make-tbn `((a . ,(make-sbf/state '((b . 1))))
                              (b . ,(make-tbf/state '((a . -1)) -1))))))
    (check-true (normalized-tbn?
                 (make-tbn `((a . ,(make-sbf/state '((a . 1) (b . -1))))
                             (b . ,(make-tbf/state '((a . -1) (b . 1)) -1))))))))

;;; Normalizes a TBN.
;;;
;;; For every TBF, removes the inputs that are not in the variables of
;;; the TBN, and adds missing inputs with 0 weight.
(define (normalize-tbn tbn)
  (define vars-0 (for/hash ([(x _) (in-hash tbn)]) (values x 0)))
  (define (normalize-tbf tbf)
    ;; Only keep the inputs which are also the variables of tbn.
    (define w-pruned (hash-intersect tbn (tbf/state-w tbf)
                                     #:combine (λ (_ w) w)))
    ;; Put in the missing inputs with weight 0.
    (define w-complete (hash-union vars-0 w-pruned #:combine (λ (_ w) w)))
    (tbf/state w-complete (tbf/state-θ tbf)))
  (for/hash ([(x tbf) (in-hash tbn)]) (values x (normalize-tbf tbf))))

(module+ test
  (test-case "normalize-tbn"
    (check-equal? (normalize-tbn
                   (hash 'a (make-sbf/state '((b . 1) (c . 3)))
                         'b (make-tbf/state '((a . -1)) -1)))
                  (hash
                   'a
                   (tbf/state '#hash((a . 0) (b . 1)) 0)
                   'b
                   (tbf/state '#hash((a . -1) (b . 0)) -1)))))

;;; Compacts (and denormalizes) a TBF by removing all inputs which
;;; are 0.
(define (compact-tbf tbf)
  (tbf/state
   (for/hash ([(k v) (in-hash (tbf/state-w tbf))]
              #:unless (zero? v))
     (values k v))
   (tbf/state-θ tbf)))

(module+ test
  (test-case "compact-tbf"
    (check-equal? (compact-tbf (tbf/state (hash 'a 0 'b 1 'c 2 'd 0) 2))
                  (tbf/state '#hash((b . 1) (c . 2)) 2))))

;;; Compacts a TBN by removing all inputs which are 0 or which are not
;;; variables of the network.
(define (compact-tbn tbn)
  (define (remove-0-non-var tbf)
    (tbf/state
     (for/hash ([(x w) (in-hash (tbf/state-w tbf))]
                #:when (hash-has-key? tbn x)
                #:unless (zero? w))
       (values x w))
     (tbf/state-θ tbf)))
  (for/hash ([(x tbf) (in-hash tbn)])
    (values x (remove-0-non-var tbf))))

(module+ test
  (test-case "compact-tbn"
    (check-equal?
     (compact-tbn (hash 'a (tbf/state (hash 'a 0 'b 1 'c 3 'd 0) 0)
                        'b (tbf/state (hash 'a -1 'b 1) -1)))
     (hash
      'a
      (tbf/state '#hash((b . 1)) 0)
      'b
      (tbf/state '#hash((a . -1) (b . 1)) -1)))))

;;; Given TBN, produces a sexp containing the description of the
;;; functions of the TBN that Org-mode can interpret as a table.
;;;
;;; Like print-org-tbfs/state, if #:headers is #f, does not print the
;;; names of the inputs of the TBFs.  If #:headers is #t, the output
;;; starts by a list giving the names of the variables, as well as the
;;; symbol 'θ to represent the column giving the thresholds of the
;;; TBF.
;;;
;;; If #:func-names is #t, the first column of the table gives the
;;; variable which the corresponding TBF updates.
;;;
;;; If both #:func-names and #:headers are #t, the first cell of the
;;; first column contains the symbol '-.
(define (print-org-tbn tbn
                       #:headers [headers #t]
                       #:func-names [func-names #t])
  (define ntbn (normalize-tbn tbn))
  (define vars-tbfs (hash-map ntbn (λ (x tbf) (cons x tbf)) #t))
  (define tbfs (map cdr vars-tbfs))
  (define tbfs-table (print-org-tbfs/state tbfs #:headers headers))
  (cond
    [(eq? func-names #t)
     (define vars (map car vars-tbfs))
     (define col-1 (if headers (cons '- vars) vars))
     (for/list ([var (in-list col-1)] [row (in-list tbfs-table)])
       (cons var row))]
    [else
     tbfs-table]))

(module+ test
  (test-case "print-org-tbn"
    (define tbn (make-tbn `((a . ,(make-sbf/state '((b . 1))))
                            (b . ,(make-tbf/state '((a . -1)) -1)))))
    (check-equal? (print-org-tbn tbn)
                  '((- a b θ) (a 0 1 0) (b -1 0 -1)))))

;;; Given an SBN, produces a sexp containing the description of the
;;; functions of the SBN that Org-mode can interpret as a table.
;;; This function is therefore very similar to print-org-tbn.
;;;
;;; Like print-org-tbfs/state, if #:headers is #f, does not print the
;;; names of the inputs of the TBFs.  If #:headers is #t, the output
;;; starts by a list giving the names of the variables.
;;;
;;; If #:func-names is #t, the first column of the table gives the
;;; variable which the corresponding TBF updates.
;;;
;;; If both #:func-names and #:headers are #t, the first cell of the
;;; first column contains the symbol '-.
(define (print-org-sbn sbn
                       #:headers [headers #t]
                       #:func-names [func-names #t])
  (define tab (print-org-tbn sbn #:headers headers #:func-names func-names))
  (define-values (tab-no-θ _) (multi-split-at
                               tab
                               (- (length (car tab)) 1)))
  tab-no-θ)

(module+ test
  (test-case "print-org-sbn"
    (define sbn (hash
                 'a
                 (tbf/state (hash 'b 2) 0)
                 'b
                 (tbf/state (hash 'a 2) 0)))
    (check-equal? (print-org-sbn sbn)
                  '((- a b) (a 0 2) (b 2 0)))))

;;; Given a TBN, constructs its interaction graph.  The nodes of this
;;; graph are labeled with pairs (variable name . threshold), while
;;; the edges are labelled with the weights.
;;;
;;; If #:zero-edges is #t, the edges with zero weights will appear in
;;; the interaction graph.
(define (tbn-interaction-graph tbn
                               #:zero-edges [zero-edges #t])
  (define ntbn (normalize-tbn tbn))
  (define ig (weighted-graph/directed
              (if zero-edges
                  (for*/list ([(tar tbf) (in-hash ntbn)]
                              [(src w) (in-hash (tbf/state-w tbf))])
                    (list w src tar))
                  (for*/list ([(tar tbf) (in-hash ntbn)]
                              [(src w) (in-hash (tbf/state-w tbf))]
                              #:unless (zero? w))
                    (list w src tar)))))
  (update-graph ig #:v-func (λ (x) (cons x (tbf/state-θ (hash-ref ntbn x))))))

(module+ test
  (test-case "tbn-interaction-graph"
    (define tbn (make-tbn `((a . ,(make-sbf/state '((b . 1))))
                            (b . ,(make-tbf/state '((a . -1)) -1)))))
    (check-equal? (graphviz (tbn-interaction-graph tbn))
                  "digraph G {\n\tnode0 [label=\"'(b . -1)\\n\"];\n\tnode1 [label=\"'(a . 0)\\n\"];\n\tsubgraph U {\n\t\tedge [dir=none];\n\t\tnode0 -> node0 [label=\"0\"];\n\t\tnode1 -> node1 [label=\"0\"];\n\t}\n\tsubgraph D {\n\t\tnode0 -> node1 [label=\"1\"];\n\t\tnode1 -> node0 [label=\"-1\"];\n\t}\n}\n")
    (check-equal? (graphviz (tbn-interaction-graph tbn #:zero-edges #f))
                  "digraph G {\n\tnode0 [label=\"'(b . -1)\\n\"];\n\tnode1 [label=\"'(a . 0)\\n\"];\n\tsubgraph U {\n\t\tedge [dir=none];\n\t}\n\tsubgraph D {\n\t\tnode0 -> node1 [label=\"1\"];\n\t\tnode1 -> node0 [label=\"-1\"];\n\t}\n}\n")))

;;; Pretty prints the node labels of the interaction graph of a TBN.
(define (pretty-print-tbn-interaction-graph ig)
  (update-graph ig #:v-func (match-lambda
                              [(cons var weight) (~a var ":" weight)])))

(module+ test
  (test-case "pretty-print-tbn-interaction-graph"
    (define tbn (make-tbn `((a . ,(make-sbf/state '((b . 1))))
                            (b . ,(make-tbf/state '((a . -1)) -1)))))
    (check-equal? (graphviz (pretty-print-tbn-interaction-graph (tbn-interaction-graph tbn)))
                  "digraph G {\n\tnode0 [label=\"b:-1\"];\n\tnode1 [label=\"a:0\"];\n\tsubgraph U {\n\t\tedge [dir=none];\n\t\tnode0 -> node0 [label=\"0\"];\n\t\tnode1 -> node1 [label=\"0\"];\n\t}\n\tsubgraph D {\n\t\tnode0 -> node1 [label=\"1\"];\n\t\tnode1 -> node0 [label=\"-1\"];\n\t}\n}\n")))

;;; Given an SBN, constructs its interaction graph.  As in
;;; tbn-interaction-graph, the nodes of this graph are labeled with
;;; the variable names, while the edges are labelled with the weights.
;;;
;;; If #:zero-edges is #t, the edges with zero weights will appear in
;;; the interaction graph.
(define (sbn-interaction-graph sbn
                               #:zero-edges [zero-edges #t])
  (update-graph (tbn-interaction-graph sbn)
                #:v-func (match-lambda
                           [(cons var _) var])))

(module+ test
  (test-case "sbn-interaction-graph"
    (define sbn (hash
                 'a
                 (tbf/state (hash 'b 2) 0)
                 'b
                 (tbf/state (hash 'a 2) 0)))
    (check-equal? (graphviz (sbn-interaction-graph sbn))
                  "digraph G {\n\tnode0 [label=\"b\"];\n\tnode1 [label=\"a\"];\n\tsubgraph U {\n\t\tedge [dir=none];\n\t\tnode0 -> node1 [label=\"2\"];\n\t\tnode0 -> node0 [label=\"0\"];\n\t\tnode1 -> node1 [label=\"0\"];\n\t}\n\tsubgraph D {\n\t}\n}\n")))