From 093a8873588d896e236ede3d728a13602b84f9cc Mon Sep 17 00:00:00 2001
From: Sam Hadow <sam.hadow@proton.me>
Date: Mon, 27 Apr 2026 09:37:36 +0200
Subject: [PATCH] variable search optimization

---
 .gitignore                  |  2 +-
 src/tea3/pretty_print.py    |  3 +++
 src/tea3/variable_search.py | 12 +++++++-----
 3 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/.gitignore b/.gitignore
index 504e2bf..ecea92e 100644
--- a/.gitignore
+++ b/.gitignore
@@ -3,4 +3,4 @@
 __pycache__/
 .pytest_cache/
 *.egg-info/
-
+*.txt
diff --git a/src/tea3/pretty_print.py b/src/tea3/pretty_print.py
index 6d68493..1e90ba6 100644
--- a/src/tea3/pretty_print.py
+++ b/src/tea3/pretty_print.py
@@ -4,6 +4,9 @@ def pretty_print(poly):
     x_terms = []
     mixed_terms = []
 
+    if isinstance(poly, int):
+        return str(poly)
+
     has_const = bool(poly.constant_coefficient())
 
     for monom in poly:
diff --git a/src/tea3/variable_search.py b/src/tea3/variable_search.py
index 0c0c8b0..6ff0fa0 100644
--- a/src/tea3/variable_search.py
+++ b/src/tea3/variable_search.py
@@ -26,12 +26,14 @@ def apply_variable_change(model, coeffs):
     return new_R
 
 def run_exhaustive(steps: int, target_reg: int = 0):
-    for idx, coeffs in enumerate(iproduct([0, 1], repeat=4)):
-        model = Tea3Model()
-        for _ in range(steps):
-            model.step()
+    model = Tea3Model()
+    for _ in range(steps):
+        model.step()
 
-        new_R = apply_variable_change(model, coeffs)
+    snapshot = model
+
+    for idx, coeffs in enumerate(iproduct([0, 1], repeat=4)):
+        new_R = apply_variable_change(snapshot, coeffs)
 
         label = "".join(map(str, coeffs))
         print(f"\n[{idx:02d}] (a1,a2,a3,a4) = {label}")