"""Operation counts for scalar multiplication on secp256k1: double-and-add, wNAF, GLV + interleaved wNAF.

Counts group additions and doublings (affine model) for random 256-bit scalars; the GLV decomposition is the
published libsecp256k1 lattice basis re-verified in secp256k1_facts.py. Writes results/baselines/scalar_mult_opcount.json.
"""
import argparse, json, random, sys, time, pathlib, statistics as st
ROOT = pathlib.Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT / "src"))
from wp34 import secp256k1 as k1
from wp34.ec import Curve
from wp34.provenance import record

E: Curve = k1.CURVE
phi = lambda P: (k1.beta * P[0] % k1.p, P[1])


def wnaf(k: int, w: int):
    naf = []
    while k > 0:
        if k & 1:
            d = k % (1 << w)
            if d >= (1 << (w - 1)):
                d -= 1 << w
            naf.append(d); k -= d
        else:
            naf.append(0)
        k >>= 1
    return naf


def mul_glv_interleaved(k: int, P, w: int = 4):
    """[k]P = [k1]P + [k2]φ(P) with |k_i| ≲ √n, interleaved wNAF (Straus/Shamir with signed windows)."""
    k1_, k2_ = k1.glv_decompose(k)
    P1, P2 = P, phi(P)
    if k1_ < 0: k1_, P1 = -k1_, E.neg(P1)
    if k2_ < 0: k2_, P2 = -k2_, E.neg(P2)
    tabs = []
    for B in (P1, P2):
        twoB = E.double(B)
        t = [B]
        for _ in range((1 << (w - 2)) - 1):
            t.append(E.add(t[-1], twoB))
        tabs.append(t)
    n1, n2 = wnaf(k1_, w), wnaf(k2_, w)
    L = max(len(n1), len(n2))
    n1 += [0] * (L - len(n1)); n2 += [0] * (L - len(n2))
    R = None
    for i in range(L - 1, -1, -1):
        R = E.double(R)
        for d, t in ((n1[i], tabs[0]), (n2[i], tabs[1])):
            if d > 0: R = E.add(R, t[(d - 1) // 2])
            elif d < 0: R = E.add(R, E.neg(t[(-d - 1) // 2]))
    return R


def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--seed", type=int, default=34); ap.add_argument("--n", type=int, default=200)
    ap.add_argument("--out", default=str(ROOT / "results/baselines/scalar_mult_opcount.json"))
    a = ap.parse_args(); rng = random.Random(a.seed); t0 = time.time()
    methods = {
        "double_and_add": lambda k, P: E.mul(k, P),
        "wnaf_w4": lambda k, P: E.mul_wnaf(k, P, 4),
        "wnaf_w5": lambda k, P: E.mul_wnaf(k, P, 5),
        "glv_wnaf_w4": lambda k, P: mul_glv_interleaved(k, P, 4),
        "glv_wnaf_w5": lambda k, P: mul_glv_interleaved(k, P, 5),
    }
    rows = {m: {"add": [], "dbl": [], "total": [], "time_s": []} for m in methods}
    for i in range(a.n):
        k = rng.randrange(1, k1.n)
        P = E.mul(rng.randrange(1, k1.n), E.G)
        ref = None
        for m, f in methods.items():
            E.ops.reset(); t = time.perf_counter(); R = f(k, P); dt = time.perf_counter() - t
            if ref is None: ref = R
            assert R == ref, (m, i)
            rows[m]["add"].append(E.ops.add); rows[m]["dbl"].append(E.ops.dbl)
            rows[m]["total"].append(E.ops.add + E.ops.dbl); rows[m]["time_s"].append(dt)
    summary = {}
    base = st.mean(rows["double_and_add"]["total"])
    for m, r in rows.items():
        summary[m] = {"add_mean": st.mean(r["add"]), "dbl_mean": st.mean(r["dbl"]), "total_mean": st.mean(r["total"]),
                      "total_sd": st.pstdev(r["total"]), "total_min": min(r["total"]), "total_max": max(r["total"]),
                      "speedup_vs_double_and_add_ops": base / st.mean(r["total"]),
                      "time_mean_ms": 1000 * st.mean(r["time_s"])}
    # GLV sub-scalar sizes
    sizes = []
    for _ in range(2000):
        k = rng.randrange(1, k1.n); a1_, a2_ = k1.glv_decompose(k)
        sizes.append((abs(a1_).bit_length(), abs(a2_).bit_length()))
    summary["glv_subscalar_bits"] = {"k1_max": max(s[0] for s in sizes), "k2_max": max(s[1] for s in sizes),
                                     "k1_mean": st.mean(s[0] for s in sizes), "k2_mean": st.mean(s[1] for s in sizes), "samples": 2000}
    record("scalar_mult_opcount_secp256k1", {"n_scalars": a.n, "model": "affine, counts group add/dbl"}, summary,
           seed=a.seed, out=a.out, started=t0)
    for m, s in summary.items():
        if m != "glv_subscalar_bits":
            print(f"{m:18s} add {s['add_mean']:7.1f} dbl {s['dbl_mean']:7.1f} total {s['total_mean']:7.1f} ± {s['total_sd']:.1f}  speedup(ops) {s['speedup_vs_double_and_add_ops']:.3f}  {s['time_mean_ms']:.1f} ms")
    print("GLV sub-scalars:", summary["glv_subscalar_bits"])


if __name__ == "__main__":
    main()