"""ECDLP study generator/driver for the Rust engine (Axis A: H04, H08, H09, H10, H11).

Studies (each writes results/synthetic_ecdlp/<study>_<tag>.jobs.json; `--run` executes locally, otherwise the jobs are
shipped to a cluster node and executed with run_batch_remote.py):
  scaling_generic  sizes 20..56, 20 random prime-order curves per size, configs bsgs / rho floyd / rho brent / rho dp /
                   rho dp+neg (r=128) / kangaroo (full interval)
  scaling_j0       sizes 24..52 on j=0 curves (p ≡ 1 mod 3), configs rho dp (r=128) / +neg / +aut3 / +neg+aut3
  sweep_r          36- and 40-bit generic curves × r ∈ {3,4,8,16,20,32,64,128} × mixed ∈ {0,2,4}, 40 instances
  parallel         48-bit generic curves, threads ∈ {1,2,4,8,16,28}, 6 instances
  tmto             bsgs memory sweep (36 bits) and rho dp-bit sweep (40 bits), 20 instances
Instances are generated with PARI/GP point counting (laptop); jobs are flat objects in the engine's batch format.
"""
import argparse, json, math, random, sys, pathlib
ROOT = pathlib.Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT / "src"))
from wp34.ec import random_prime_order_curve, cube_roots_of_unity, endomorphism_lambda

OUT = ROOT / "results/synthetic_ecdlp"


def gen_instances(sizes, per_size, seed, j0):
    rng = random.Random(seed); inst = []
    for bits in sizes:
        for i in range(per_size):
            E = random_prime_order_curve(bits, rng, j0=j0)
            d = rng.randrange(1, E.n); Q = E.mul(d, E.G)
            rec = {"bits": bits, "idx": i, "p": str(E.p), "a": str(E.a), "b": str(E.b), "n": str(E.n),
                   "gx": str(E.G[0]), "gy": str(E.G[1]), "qx": str(Q[0]), "qy": str(Q[1]), "expected_d": str(d)}
            if j0:
                for beta in cube_roots_of_unity(E.p):
                    try:
                        lam = endomorphism_lambda(E, beta); rec["beta"] = str(beta); rec["lambda"] = str(lam); break
                    except ValueError:
                        continue
                assert "lambda" in rec
            inst.append(rec)
    return inst


def base(rec):
    return {k: rec[k] for k in ("p", "a", "b", "n", "gx", "gy", "qx", "qy", "expected_d")}


def dp_default(bits):
    return max(1, bits // 4 - 2)


def study_scaling_generic(seed):
    sizes = [20, 24, 28, 32, 36, 40, 44, 48, 52, 56]
    inst = gen_instances(sizes, 60, seed, False); jobs = []
    for rec in inst:
        b = base(rec); lab = f"g{rec['bits']}b-{rec['idx']}"; bits = rec["bits"]
        if bits <= 48:
            jobs.append(dict(b, alg="bsgs", seed=seed, threads=1, label=lab + "/bsgs"))
        if bits <= 48:
            jobs.append(dict(b, alg="rho", cycle="floyd", r=20, seed=seed, threads=1, label=lab + "/rho_floyd"))
            jobs.append(dict(b, alg="rho", cycle="brent", r=20, seed=seed, threads=1, label=lab + "/rho_brent"))
        jobs.append(dict(b, alg="rho", cycle="dp", r=20, dp=dp_default(bits), seed=seed, threads=1, label=lab + "/rho_dp"))
        jobs.append(dict(b, alg="rho", cycle="dp", r=128, dp=dp_default(bits), seed=seed, threads=1, label=lab + "/rho_dp_r128"))
        jobs.append(dict(b, alg="rho", cycle="dp", r=128, dp=dp_default(bits), neg=True, seed=seed, threads=1, label=lab + "/rho_dp_neg"))
        jobs.append(dict(b, alg="kangaroo", dp=dp_default(bits), seed=seed, threads=1, label=lab + "/kangaroo"))
    return jobs, {"sizes": sizes, "per_size": 60}


def study_scaling_j0(seed):
    sizes = [24, 28, 32, 36, 40, 44, 48, 52]
    inst = gen_instances(sizes, 60, seed + 1, True); jobs = []
    for rec in inst:
        b = base(rec); lab = f"j{rec['bits']}b-{rec['idx']}"; bits = rec["bits"]; dp = dp_default(bits)
        jobs.append(dict(b, alg="rho", cycle="dp", r=128, dp=dp, seed=seed, threads=1, label=lab + "/rho_dp"))
        jobs.append(dict(b, alg="rho", cycle="dp", r=128, dp=dp, neg=True, seed=seed, threads=1, label=lab + "/rho_dp_neg"))
        jobs.append(dict(b, alg="rho", cycle="dp", r=128, dp=dp, aut3=True, beta=rec["beta"], seed=seed, threads=1, label=lab + "/rho_dp_aut3", **{"lambda": rec["lambda"]}))
        jobs.append(dict(b, alg="rho", cycle="dp", r=128, dp=dp, neg=True, aut3=True, beta=rec["beta"], seed=seed, threads=1, label=lab + "/rho_dp_neg_aut3", **{"lambda": rec["lambda"]}))
        jobs.append(dict(b, alg="rho", cycle="dp", r=20, dp=dp, neg=True, aut3=True, beta=rec["beta"], seed=seed, threads=1, label=lab + "/rho_dp_neg_aut3_r20", **{"lambda": rec["lambda"]}))
    return jobs, {"sizes": sizes, "per_size": 60}


def study_sweep_r(seed):
    inst = gen_instances([36, 40], 40, seed + 2, False); jobs = []
    for rec in inst:
        b = base(rec); lab = f"g{rec['bits']}b-{rec['idx']}"; dp = dp_default(rec["bits"])
        for r in (3, 4, 8, 16, 20, 32, 64, 128):
            for mixed in (0, 2, 4):
                kw = {"mixed": mixed} if mixed else {}
                jobs.append(dict(b, alg="rho", cycle="dp", r=r, dp=dp, seed=seed, threads=1, label=f"{lab}/r{r}_m{mixed}", **kw))
        # negative control: a bad walk — r = 2 (far from random), and Pollard's original 3-partition analogue r=3 is above
        jobs.append(dict(b, alg="rho", cycle="dp", r=2, dp=dp, seed=seed, threads=1, label=f"{lab}/r2_m0"))
    return jobs, {"sizes": [36, 40], "per_size": 40, "r": [2, 3, 4, 8, 16, 20, 32, 64, 128], "mixed": [0, 2, 4]}


def study_parallel(seed):
    inst = gen_instances([56], 12, seed + 3, False); jobs = []
    for rec in inst:
        b = base(rec); lab = f"g56b-{rec['idx']}"
        for t in (1, 2, 4, 8, 16, 28):
            jobs.append(dict(b, alg="rho", cycle="dp", r=128, dp=12, neg=True, seed=seed, threads=t, label=f"{lab}/t{t}"))
    return jobs, {"sizes": [56], "per_size": 12, "threads": [1, 2, 4, 8, 16, 28]}


def study_tmto(seed):
    jobs = []
    for rec in gen_instances([36], 20, seed + 4, False):
        b = base(rec); lab = f"g36b-{rec['idx']}"
        for mem in (2**10, 2**12, 2**14, 2**16, 2**18, 2**20):
            jobs.append(dict(b, alg="bsgs", memory=str(mem), seed=seed, threads=1, label=f"{lab}/bsgs_m{mem}"))
    for rec in gen_instances([40], 20, seed + 5, False):
        b = base(rec); lab = f"g40b-{rec['idx']}"
        for dp in (2, 4, 6, 8, 10, 12, 14):
            jobs.append(dict(b, alg="rho", cycle="dp", r=128, dp=dp, seed=seed, threads=1, label=f"{lab}/rho_dp{dp}"))
    return jobs, {"bsgs_bits": 36, "rho_bits": 40, "per_size": 20}


STUDIES = {"scaling_generic": study_scaling_generic, "scaling_j0": study_scaling_j0, "sweep_r": study_sweep_r,
           "parallel": study_parallel, "tmto": study_tmto}


def main():
    ap = argparse.ArgumentParser(); ap.add_argument("study", choices=list(STUDIES) + ["all"]); ap.add_argument("--seed", type=int, default=34)
    a = ap.parse_args(); OUT.mkdir(parents=True, exist_ok=True)
    for name in (STUDIES if a.study == "all" else [a.study]):
        jobs, params = STUDIES[name](a.seed)
        (OUT / f"{name}.jobs.json").write_text(json.dumps(jobs))
        (OUT / f"{name}.params.json").write_text(json.dumps(dict(params, seed=a.seed, study=name)))
        print(f"{name}: {len(jobs)} jobs → {OUT / (name + '.jobs.json')}")


if __name__ == "__main__":
    main()