stronger abstraction

.gitignore

@@ -3,4 +3,4 @@
 __pycache__/
 .pytest_cache/
 *.egg-info/
-*.txt
+

src/tea3/cli.py

@@ -1,6 +1,5 @@
 from tea3.pretty_print import pretty_print
 from tea3.tea3model import Tea3Model
-from tea3.variable_search import run_exhaustive
 
 
 def prompt_int(message: str, lo: int, hi: int) -> int:
@@ -16,80 +15,30 @@ def prompt_int(message: str, lo: int, hi: int) -> int:
         print(f"Value must be between {lo} and {hi} (inclusive).")
 
 
-def prompt_choice(message: str, choices: set[int]) -> int:
-    choices_str = ", ".join(map(str, sorted(choices)))
-    while True:
-        raw = input(message).strip()
-        try:
-            value = int(raw)
-        except ValueError:
-            print("Please enter a number.")
-            continue
-        if value in choices:
-            return value
-        print(f"Value must be one of: {choices_str}")
+def main() -> None:
+    print("=" * 50)
+    print(" Tea3 Model ")
+    print("=" * 50)
 
+    steps = prompt_int("\nHow many steps do you want to run? (1–100): ", 1, 100)
 
-def run_classic_cli():
     print("\nR registers are indexed 0–7; bits within each register are 0–7.")
-    print("Enter -1 to print all registers, or all bits.")
-
-    steps = prompt_int("How many steps do you want to run? (1–100): ", 1, 100)
-    reg = prompt_int("Which R register do you want to inspect? (-1 or 0–7): ", -1, 7)
-    bit = prompt_int("Which bit of that register?              (-1 or 0–7): ", -1, 7)
+    reg = prompt_int("Which R register do you want to inspect? (0–7): ", 0, 7)
+    bit = prompt_int("Which bit of that register?              (0–7): ", 0, 7)
 
+    print(f"\nRunning {steps} step(s), watching R[{reg}][{bit}]")
     print("-" * 50)
 
     model = Tea3Model()
 
     for i in range(steps):
         model.step()
-        print(f"\n[Step {i + 1}]")
-
-        regs = range(8) if reg == -1 else [reg]
-        bits = range(8) if bit == -1 else [bit]
-
-        for r in regs:
-            for b in bits:
-                poly = model.R_bits[r][b]
-                print(f"R_bits[{r}][{b}] =")
+        poly = model.R_bits[reg][bit]
+        print(f"\n[Step {i+1}]  R_bits[{reg}][{bit}] =")
         print(pretty_print(poly))
-                print()
 
     print("\n" + "=" * 50)
     print("Done.")
 
-def run_exhaustive_cli():
-    print("\nR registers are indexed 0–7; bits within each register are 0–7.")
-    print("Enter -1 to print all bits in the chosen register.")
-
-    steps = prompt_int("How many steps? (1–100): ", 1, 100)
-    target_reg = prompt_int("Target register (0–7): ", 0, 7)
-    target_bit = prompt_int("Target bit (-1 or 0–7): ", -1, 7)
-
-    print("-" * 50)
-    run_exhaustive(steps, target_reg, target_bit)
-
-    print("\n" + "=" * 50)
-    print("Done.")
-
-
-
-def main():
-    print("=" * 50)
-    print(" Tea3 Model ")
-    print("=" * 50)
-
-    print("\nChoose a mode:")
-    print("  1) Classic inspection")
-    print("  2) Exhaustive variable-change search")
-
-    mode = prompt_choice("Your choice (1 or 2): ", {1, 2})
-
-    if mode == 1:
-        run_classic_cli()
-    else:
-        run_exhaustive_cli()
-
 
 main()

src/tea3/pretty_print.py

@@ -1,12 +1,11 @@
 def pretty_print(poly):
+    ring = poly.parent()
+
     R_terms = []
     r_terms = []
     x_terms = []
     mixed_terms = []
 
-    if isinstance(poly, int):
-        return str(poly)
-
     has_const = bool(poly.constant_coefficient())
 
     for monom in poly:
@@ -29,11 +28,10 @@ def pretty_print(poly):
 
     parts = []
 
-    if has_const:
-        parts.append("1")
-
-    if R_terms:
+    if R_terms or (has_const and not r_terms):
         parts.append("f(Ri)")
+    elif has_const:
+        parts.append("1")
 
     if r_terms:
         parts.append("f(ri)")

src/tea3/tea3model.py

@@ -12,24 +12,24 @@ class Tea3Model:
         self.step_count = 0
 
         names = (
-            [f"x{i}{j}" for i in range(5) for j in range(8)] +  # 0–39
-            [f"y{i}{j}" for i in range(5) for j in range(8)] +  # 40–79
-            [f"r{i}{j}" for i in range(5) for j in range(8)] +  # 80–119
-            [f"R{i}{j}" for i in range(8) for j in range(8)] +  # 120–183
-            ["g"]                                               # 184
+            [f"x{i}{j}" for i in range(5) for j in range(8)] +
+            [f"r{i}{j}" for i in range(5) for j in range(8)] +
+            [f"R{i}{j}" for i in range(8) for j in range(8)] +
+            ["g"]
         )
 
         name_string = ",".join(names)
         self.S = BooleanPolynomialRing(len(names), name_string)
         self.v = self.S.gens()
 
-        self.x_bits = [list(self.v[i*8       : (i+1)*8      ]) for i in range(5)]
-        self.y_bits = [list(self.v[40 + i*8  : 40 + (i+1)*8]) for i in range(5)]
-        self.r_bits = [list(self.v[80 + i*8  : 80 + (i+1)*8]) for i in range(5)]
-        self.R_bits = [list(self.v[120 + i*8 : 120 + (i+1)*8]) for i in range(8)]
+        self.x_bits = [list(self.v[i*8:(i+1)*8])             for i in range(5)]
+        self.r_bits = [list(self.v[40 + i*8 : 40 + (i+1)*8]) for i in range(5)]
+        self.R_bits = [list(self.v[80 + i*8 : 80 + (i+1)*8]) for i in range(8)]
+
         self.g = self.v[-1]
 
     def _abstract_R(self):
+        s    = self.step_count
         one  = self.S.one()
         zero = self.S.zero()
 
@@ -38,7 +38,7 @@ class Tea3Model:
                 poly = self.R_bits[i][j]
 
                 groups  = {}
-                pure_xyr = zero
+                pure_xr = zero
                 const   = one if bool(poly.constant_coefficient()) else zero
 
                 for monom in poly:
@@ -46,20 +46,20 @@ class Tea3Model:
                     if not term_vars:
                         continue
 
-                    xyr_vars = [v for v in term_vars if str(v)[0] in ('x', 'y', 'r')]
-                    Rg_vars  = [v for v in term_vars if str(v)[0] in ('R', 'g')]
+                    xr_vars = [v for v in term_vars if str(v)[0] in ('x', 'r')]
+                    Rf_vars = [v for v in term_vars if str(v)[0] in ('R', 'f', 'g')]
 
-                    xyr_mono = reduce(mul, (self.S(v) for v in xyr_vars), one)
-                    xyr_key  = frozenset(str(v) for v in xyr_vars)
+                    xr_mono = reduce(mul, (self.S(v) for v in xr_vars), one)
+                    xr_key  = frozenset(str(v) for v in xr_vars)
 
-                    if not Rg_vars:
-                        pure_xyr += xyr_mono
+                    if not Rf_vars:
+                        pure_xr += xr_mono
                     else:
-                        groups[xyr_key] = xyr_mono
+                        groups[xr_key] = xr_mono   # Rf_sum is irrelevant now
 
-                result = pure_xyr + const
-                for xyr_key, xyr_mono in groups.items():
-                    result += xyr_mono * self.g
+                result = pure_xr + const
+                for xr_key, xr_mono in groups.items():
+                    result += xr_mono * self.g      # same g every time
 
                 self.R_bits[i][j] = result
 
@@ -101,7 +101,6 @@ class Tea3Model:
 
         return R7
 
-
 def S(r):
     # placeholder
     return r

src/tea3/variable_search.py

@@ -1,50 +0,0 @@
-from itertools import product as iproduct
-from tea3.tea3model import Tea3Model
-from tea3.pretty_print import pretty_print
-
-def apply_variable_change(model, coeffs):
-    orig_x = [list(model.v[i*8 : (i+1)*8]) for i in range(5)]
-    orig_y = model.y_bits
-
-    subs = {}
-
-    for i in range(1, 5):
-        for j in range(8):
-            subs[orig_x[i][j]] = orig_y[i][j]
-
-    for j in range(8):
-        new_expr = orig_y[0][j]
-        for i, ai in enumerate(coeffs, start=1):
-            if ai:
-                new_expr = new_expr + orig_y[i][j]
-        subs[orig_x[0][j]] = new_expr
-
-    new_R = []
-    for i in range(8):
-        row = [poly.subs(subs) for poly in model.R_bits[i]]
-        new_R.append(row)
-    return new_R
-
-def run_exhaustive(steps: int, target_reg: int = 0, target_bit: int = -1):
-    model = Tea3Model()
-    for _ in range(steps):
-        model.step()
-
-    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}")
-
-        if target_bit == -1:
-            bits = range(8)
-        else:
-            if not (0 <= target_bit < 8):
-                raise ValueError("target_bit must be in [0, 7] or -1")
-            bits = [target_bit]
-
-        for j in bits:
-            poly = new_R[target_reg][j]
-            print(f"  R[{target_reg}][{j}] = {pretty_print(poly)}")