From e9ecbd8a7c8b92d2a5564daac031c57581c08927 Mon Sep 17 00:00:00 2001
From: Sergiu Ivanov <sivanov@colimite.fr>
Date: Tue, 8 Aug 2023 09:21:12 +0200
Subject: [PATCH] Type pretty-print-tbn-interaction-graph.

---
 scribblings/tbn.scrbl | 11 +++++++++++
 tbn.rkt               | 15 ++++++++++++++-
 2 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/scribblings/tbn.scrbl b/scribblings/tbn.scrbl
index e1fe23f..a8959f7 100644
--- a/scribblings/tbn.scrbl
+++ b/scribblings/tbn.scrbl
@@ -436,6 +436,17 @@ will also appear in the interaction graph.
                               #:zero-edges #f))
 ]}
 
+@defproc[(pretty-print-tbn-interaction-graph [ig Graph]) Graph]{
+
+Pretty prints the node labels of the interaction graph of a TBN.
+
+@ex[
+(dotit (pretty-print-tbn-interaction-graph
+        (tbn-interaction-graph (hash 'a (tbf/state (hash 'b 1) 0)
+                                     'b (tbf/state (hash 'a -1) -1)))))
+]}
+
+
 @section{Reading and printing TBNs and SBNs}
 
 @defproc[(parse-org-tbn [tab (Listof (Listof (U Symbol Real)))]
diff --git a/tbn.rkt b/tbn.rkt
index d7728c6..b10367b 100644
--- a/tbn.rkt
+++ b/tbn.rkt
@@ -48,7 +48,7 @@
 
    TBN sbn? tbn->network
    build-tbn-state-graph normalized-tbn? normalize-tbn compact-tbn
-   tbn-interaction-graph
+   tbn-interaction-graph pretty-print-tbn-interaction-graph
    parse-org-tbn read-org-tbn read-org-sbn tbn->lists sbn->lists
    )
 
@@ -524,6 +524,19 @@
       (check-equal? (graphviz (tbn-interaction-graph tbn #:zero-edges #f))
                     "digraph G {\n\tnode0 [label=\"'(b . -1)\"];\n\tnode1 [label=\"'(a . 0)\"];\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-print-tbn-interaction-graph (-> Graph Graph))
+  (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 (hash 'a (tbf/state (hash 'b 1) 0)
+                        'b (tbf/state (hash 'a -1) -1)))
+      (check-equal? (graphviz (pretty-print-tbn-interaction-graph
+                               (tbn-interaction-graph tbn)))
+                    "digraph G {\n\tnode0 [label=\"a:0\"];\n\tnode1 [label=\"b:-1\"];\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")))
+
   (: parse-org-tbn (->* ((Listof (Listof (U Symbol Real))))
                         (#:headers Boolean
                          #:func-names Boolean)