From 825220c71130cc5db29aca31e0f3c4a07f467939 Mon Sep 17 00:00:00 2001
From: Ben <ben.curtis@zama.ai>
Date: Mon, 5 Feb 2024 15:23:03 +0000
Subject: [PATCH] push WIP

---
 estimator/lwe_primal.py | 21 ++++++++++++++++++---
 1 file changed, 18 insertions(+), 3 deletions(-)

diff --git a/estimator/lwe_primal.py b/estimator/lwe_primal.py
index 22deeeb..7598760 100644
--- a/estimator/lwe_primal.py
+++ b/estimator/lwe_primal.py
@@ -300,11 +300,13 @@ def gaussian_heuristic_log_input(r):
 
         d = len(r)
         r = [log(x) for x in r]
+        
         for i, _ in enumerate(r):
-            if gaussian_heuristic_log_input(r[i:]) < D.stddev**2 * (d - i):
-                return ZZ(d - (i - 1))
+            if gaussian_heuristic_log_input(r[d - (i + 1):]) > D.stddev**2 * (i + 1):
+                return ZZ(i)
         return ZZ(2)
 
+
     @staticmethod
     @cached_function
     def cost(
@@ -588,8 +590,20 @@ def __call__(
             log_level=log_level + 1,
         )
 
+
+        def find_zeta_max(params, red_cost_model):
+            usvp_cost = primal_usvp(params, red_cost_model = red_cost_model)["rop"]
+            zeta_max = 1
+            while zeta_max < params.n:
+                if params.Xs.support_size(zeta_max) > usvp_cost:
+                    # double it for mitm
+                    return 2 * zeta_max
+                zeta_max +=1
+            return params.n
+
         if zeta is None:
-            with local_minimum(0, params.n, log_level=log_level) as it:
+            zeta_max = find_zeta_max(params, red_cost_model)
+            with local_minimum(0, zeta_max, log_level=log_level) as it:
                 for zeta in it:
                     it.update(
                         f(
@@ -603,6 +617,7 @@ def __call__(
         else:
             cost = f(zeta=zeta)
 
+
         cost["tag"] = tag
         cost["problem"] = params