"""Discrete logarithms with preprocessing (Bernstein–Lange 2013 "non-uniform cracks"; Corrigan-Gibbs–Kogan 2018 bound).

Offline: T distinguished points are produced by random walks of mean length W = ⌈√(n/T)⌉ (DP mask with ≈ log2 W bits)
and stored with their logarithms (walks start at known multiples of G, so each DP has a known log).
Online: a walk from Q (random start a·G + Q) runs until it reaches a stored DP, at which point d is recovered.
Prediction: online cost ≈ n/(T·W) ≈ √(n/T) group operations; total (T·W + online) is never below √n but the
*per-instance* cost drops as 1/√T — the "non-uniform" speedup. ST² = Ω̃(n) (Corrigan-Gibbs–Kogan) bounds this.
Pure-Python reference on 2^22–2^26 groups; the Rust engine repeats this at larger sizes.
"""
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
from wp34.dlp import RAddingWalk
from wp34.provenance import record


def experiment(bits, T, instances, rng, r=20):
    E = random_prime_order_curve(bits, rng)
    n, G = E.n, E.G
    # walk independent of Q: M_j = u_j G (v_j = 0) so the logs of all visited points are known offline
    walk = RAddingWalk(E, G, G, n, r, rng)
    walk.v = [0] * r
    walk.M = [E.mul(u, G) for u in walk.u]
    W = max(2, math.isqrt(n // T))
    dp_bits = max(1, round(math.log2(W)))
    mask = (1 << dp_bits) - 1
    E.ops.reset()
    table = {}
    offline_steps = 0
    max_len = 20 * (1 << dp_bits)
    while len(table) < T:
        a = rng.randrange(1, n); P = E.mul(a, G); L = 0
        while P is not None and (P[0] & mask) != 0 and L < max_len:
            P, a, _ = walk.step(P, a, 0); L += 1
        offline_steps += L
        if P is not None and (P[0] & mask) == 0 and P not in table:
            table[P] = a
    offline_ops = E.ops.group_ops()
    online = []
    for _ in range(instances):
        d = rng.randrange(1, n); Q = E.mul(d, G)
        E.ops.reset()
        found = None
        while found is None:
            a = rng.randrange(1, n); P = E.add(E.mul(a, G), Q); L = 0  # P = (a + d) G
            E.ops.reset()
            while True:
                if P is None:
                    break
                if (P[0] & mask) == 0:
                    if P in table:
                        found = (table[P] - a) % n
                    break
                P, a, _ = walk.step(P, a, 0); L += 1
                if L > 50 * max_len:
                    break
            online.append(E.ops.group_ops()) if found is not None else None
        assert found == d
    return {"bits": bits, "n": str(n), "T": T, "W": W, "dp_bits": dp_bits, "offline_group_ops": offline_ops,
            "offline_steps": offline_steps, "table_entries": len(table), "instances": instances,
            "online_mean": st.mean(online), "online_median": st.median(online), "online_sd": st.pstdev(online),
            "online_ci95_halfwidth": 1.96 * st.pstdev(online) / math.sqrt(len(online)),
            "pred_sqrt_n_over_T": math.sqrt(n / T), "sqrt_n": math.sqrt(n),
            "ratio_online_to_pred": st.mean(online) / math.sqrt(n / T),
            "per_instance_speedup_vs_sqrt_n": math.sqrt(n) / st.mean(online)}


def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--seed", type=int, default=34)
    ap.add_argument("--out", default=str(ROOT / "results/synthetic_ecdlp/preprocessing_python.json"))
    a = ap.parse_args(); rng = random.Random(a.seed); t0 = time.time(); rows = []
    for bits in (22, 26):
        for T in (2**4, 2**6, 2**8, 2**10, 2**12):
            r = experiment(bits, T, 60, rng); rows.append(r)
            print(f"{bits}b T={T:5d} W={r['W']:5d}  offline {r['offline_group_ops']:9d}  online {r['online_mean']:9.1f} ± {r['online_ci95_halfwidth']:.1f}  pred √(n/T) {r['pred_sqrt_n_over_T']:8.1f}  ratio {r['ratio_online_to_pred']:.2f}  speedup/inst {r['per_instance_speedup_vs_sqrt_n']:.1f}", flush=True)
    record("dlog_with_preprocessing_python", {"r": 20, "instances": 60}, {"rows": rows}, seed=a.seed, out=a.out, started=t0)


if __name__ == "__main__":
    main()