From 23282026e871668d789d5c98e51f13afd7e2a9e2 Mon Sep 17 00:00:00 2001
From: Alexander Varga <thealexvarga@gmail.com>
Date: Thu, 26 Sep 2024 15:08:01 -0400
Subject: [PATCH] generate_image: add --close and --filter flags

---
 scripts/generate_image.py | 31 ++++++++++++++++++++++++-------
 1 file changed, 24 insertions(+), 7 deletions(-)

diff --git a/scripts/generate_image.py b/scripts/generate_image.py
index 5383128e..095ef967 100644
--- a/scripts/generate_image.py
+++ b/scripts/generate_image.py
@@ -4,7 +4,8 @@
 Example usage:
 
 ```sh
-$ python scripts/generate_image.py equational_theories/Basic.lean
+$ pip install pillow
+$ python scripts/generate_image.py equational_theories/Basic.lean --close --filter
 $ open equational_theories/Basic.png
 ```
 """
@@ -21,6 +22,19 @@ def name_to_ind(name):
 KNOWN_IMPLIES_COLOR = (0, 255, 0)
 KNOWN_NOT_IMPLIES_COLOR = (255, 0, 0)
 
+def close(known_implies):
+    # copilot wrote this
+    new_pairs = closure = set(known_implies)
+    while new_pairs:
+        new_pairs = {
+            (a, d)
+            for a, b in new_pairs
+            for c, d in known_implies
+            if b == c
+        } - closure
+        closure |= new_pairs
+    return closure
+
 if __name__ == '__main__':
     try:
         file_name = argv[1]
@@ -29,16 +43,19 @@ def name_to_ind(name):
         print('Usage: python process_implications.py <file_name.lean>')
         exit(1)
 
-
     universe, known_implies, known_not_implies = parse_proofs_file(file_name)
-    # all_unknown = get_unknown_implications(universe, known_implies, known_not_implies)
 
-    inds = [name_to_ind(e) for e in universe]
-    min_ind = min(inds)
-    max_ind = max(inds)
+    if '--close' in argv:
+        known_implies = close(known_implies)
+
+    inds = {e: name_to_ind(e) for e in universe}
+    if '--filter' in argv:
+        inds = {kv[0]: i for i, kv in enumerate(sorted(inds.items(), key=lambda kv: kv[1]))}
+    min_ind = min(inds.values())
+    max_ind = max(inds.values())
     size = max_ind - min_ind + 1
     def ind(name):
-        return name_to_ind(name) - min_ind
+        return inds[name] - min_ind
 
     data = [UNKNOWN_COLOR] * (size * size)
     for known in known_implies: