"""Index-calculus applicability probe on prime-field curves (Axis A, hypothesis on Semaev decomposition).

Question: with a factor base F = {P ∈ E(F_p) : x(P) < B}, what fraction of random points R decomposes as R = P1 + P2
(P_i ∈ F), and what is the implied cost of a relation versus the generic √n cost?
Theory (Semaev 2004; Gaudry 2009; Diem 2011): over a prime field there is no subfield structure; a random point is a
sum of m factor-base points with probability ≈ |F|^m / (m!·n), so with |F| = n^{1/m}·c the relation-collection cost
is ≈ m!·n^{1−1/m}·(cost of decomposition), never below n^{1/2} for m = 2 and worse for m ≥ 3 once the decomposition
cost (solving S_{m+1} = 0) is included. We measure the m = 2 decomposition rate empirically on 16–28-bit curves.
"""
import argparse, math, random, sys, time, pathlib, statistics as st
ROOT = pathlib.Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT / "src"))
from wp34.ec import random_prime_order_curve, sqrt_mod
from wp34.provenance import record


def run(bits, frac_exp, trials, rng):
    E = random_prime_order_curve(bits, rng)
    n, p = E.n, E.p
    B = int(p ** frac_exp)          # factor base: x < B, |F| ≈ B (two points per x on average… half the x are on-curve ⇒ ≈ B)
    # enumerate the factor base (points with x < B); store by x for O(1) membership
    F = {}
    for x in range(B):
        rhs = (x * x * x + E.a * x + E.b) % p
        y = sqrt_mod(rhs, p)
        if y is not None:
            F[x] = y
    size = sum(1 if y == 0 else 2 for y in F.values())
    hits = 0
    dec_cost = 0
    for _ in range(trials):
        R = E.random_point(rng)
        # R = P1 + P2 ⇔ R − P1 ∈ F for some P1 ∈ F: try all P1 (|F| subtractions) — the naive decomposition
        found = False
        for x, y in F.items():
            for s in (y, p - y) if y else (0,):
                dec_cost += 1
                S = E.sub(R, (x, s))
                if S is None or S[0] < B:
                    found = True
                    break
            if found:
                break
        hits += found
    rate = hits / trials
    pred = 1 - math.exp(-size * size / (2 * n))     # m = 2 prediction: P[R ∈ F + F] ≈ 1 − exp(−|F|²/(2n))
    return {"bits": bits, "p": str(p), "n": str(n), "B": B, "factor_base_size": size, "trials": trials, "hits": hits,
            "decomposition_rate": rate, "predicted_rate": pred,
            "relations_needed": size, "expected_cost_group_ops": (size / rate) * size if rate else None,
            "sqrt_n": math.sqrt(n), "cost_over_sqrt_n": ((size / rate) * size / math.sqrt(n)) if rate else None,
            # even with a FREE decomposition oracle: |F| relations at 1/rate trials each, then sparse linear algebra ~ |F|²
            "oracle_relation_cost": (size / rate) if rate else None,
            "oracle_relation_cost_over_sqrt_n": (size / rate / math.sqrt(n)) if rate else None,
            "linear_algebra_cost_over_sqrt_n": size * size / math.sqrt(n)}


def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--seed", type=int, default=34)
    ap.add_argument("--out", default=str(ROOT / "results/synthetic_ecdlp/index_calculus_probe.json"))
    a = ap.parse_args(); rng = random.Random(a.seed); t0 = time.time()
    rows = []
    for bits, trials in [(16, 400), (20, 300), (24, 200), (28, 120)]:
        for fe in (0.5, 0.6):
            if bits * fe > 15.5:  # keep |F| ≤ ~50k so the naive decomposition stays affordable
                continue
            r = run(bits, fe, trials, rng); r["frac_exp"] = fe; rows.append(r)
            print(f"{bits}b  |F|=p^{fe}: size {r['factor_base_size']:6d} rate {r['decomposition_rate']:.4f} (pred {r['predicted_rate']:.4f})  naive cost/√n ≈ {round(r['cost_over_sqrt_n'],1)}  oracle-relations/√n ≈ {round(r['oracle_relation_cost_over_sqrt_n'],2)}  linalg/√n ≈ {round(r['linear_algebra_cost_over_sqrt_n'],1)}")
    record("index_calculus_decomposition_probe", {"m": 2, "sizes": [16, 20, 24, 28]}, {"rows": rows}, seed=a.seed, out=a.out, started=t0)


if __name__ == "__main__":
    main()