From bf3146cf98d16534c1d6a883e83d29fbd3202b6a Mon Sep 17 00:00:00 2001
From: Sam Hadow <sam.hadow@proton.me>
Date: Tue, 7 Jul 2026 13:23:46 +0200
Subject: [PATCH] precompute model

---
 .gitignore             |  1 +
 src/tea3/cli.py        | 11 +++++--
 src/tea3/precompute.py | 70 ++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 80 insertions(+), 2 deletions(-)
 create mode 100644 src/tea3/precompute.py

diff --git a/.gitignore b/.gitignore
index ecea92e..9a9ac17 100644
--- a/.gitignore
+++ b/.gitignore
@@ -4,3 +4,4 @@ __pycache__/
 .pytest_cache/
 *.egg-info/
 *.txt
+*.sobj
diff --git a/src/tea3/cli.py b/src/tea3/cli.py
index 96673d8..75f05a7 100644
--- a/src/tea3/cli.py
+++ b/src/tea3/cli.py
@@ -7,6 +7,7 @@ from tea3.sbox import run_sbox
 from tea3.variable_xor import run_variable_xor, run_exhaustive_xor
 from tea3.f31f32 import run_f31f32
 from tea3.cube_attack import run_cube_attack
+from tea3.precompute import run_precompute_cli, run_load_cli
 
 
 def run_classic_cli():
@@ -214,11 +215,13 @@ def main():
         print("  6) F31, F32 analysis")
         print("  7) Exhaustive XOR search")
         print("  8) Cube attack search")
+        print("  9) Precompute and save model")
+        print("  10) Load model")
         print("  0) Exit")
 
         mode = prompt_choice(
-            "Your choice (0-8): ",
-            {0, 1, 2, 3, 4, 5, 6, 7, 8}
+            "Your choice (0-10): ",
+            {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10}
         )
 
         if mode == 0:
@@ -240,6 +243,10 @@ def main():
             run_exhaustive_xor_cli()
         elif mode == 8:
             run_cube_attack_cli()
+        elif mode == 9:
+            run_precompute_cli()
+        elif mode == 10:
+            model = run_load_cli()
 
         input("\nPress Enter to return to the main menu...")
         print()
diff --git a/src/tea3/precompute.py b/src/tea3/precompute.py
new file mode 100644
index 0000000..5c65299
--- /dev/null
+++ b/src/tea3/precompute.py
@@ -0,0 +1,70 @@
+import os
+from tea3.tea3model import Tea3Model
+from tea3.utils import set_bits
+from tea3.cliutils import prompt_int, prompt_choice
+from sage.all import save, load
+
+def run_precompute_cli():
+    """
+    CLI interface:
+    build a model, optionally fix bits, run n steps (with or without abstraction), and save the model to a file.
+    """
+    print("\nPrecompute and save model")
+    print("R registers 0–7; x,r registers 0–4; bits 0–7.")
+
+    steps = prompt_int("How many steps to precompute? (1–100): ", 1, 100)
+
+    raw0 = input("Bits to set to 0: ").strip()
+    zero_bits = raw0.split() if raw0 else []
+
+    raw1 = input("Bits to set to 1: ").strip()
+    one_bits = raw1.split() if raw1 else []
+
+    skip_abs = prompt_choice("Skip abstraction (R-register monomial replacement)? (0/1): ", {0, 1})
+
+    filename = input("File to save the model (e.g. 'model_8steps.sobj'): ").strip()
+    if not filename:
+        print("No file name given. Aborting.")
+        return
+
+    # build model
+    model = Tea3Model()
+
+    try:
+        set_bits(model, zero_bits, 0)
+        set_bits(model, one_bits, 1)
+    except ValueError as e:
+        print(f"Error setting bits: {e}")
+        return
+
+    print(f"Running {steps} steps")
+    for i in range(steps):
+        model.step(skip_abstract=skip_abs)
+        print(f"  Step {i+1} done.")
+
+    # save
+    try:
+        save(model, filename)
+        print(f"Model saved to '{os.path.abspath(filename)}'")
+    except Exception as e:
+        print(f"Failed to save model: {e}")
+
+def run_load_cli():
+    """
+    CLI interface:
+    Ask for a saved model file path, load it, and return the model.
+    Prints the number of steps already computed.
+    Returns None on error.
+    """
+    filename = input("Enter the path to the saved model file: ").strip()
+    if not filename:
+        print("No file name given.")
+        return None
+
+    try:
+        model = load(filename)
+        print(f"Model loaded. Precomputed steps: {model.step_count}")
+        return model
+    except Exception as e:
+        print(f"Failed to load model: {e}")
+        return None