"""Reference (pure Python, exact) generic discrete-logarithm algorithms on an elliptic-curve group.

These are the *validation* implementations: slow but transparent. The Rust engine in
experiments/synthetic_ecdlp/rho_engine is the benchmarking implementation and must agree with these.

All functions take a Curve E with prime order n (E.n), a generator G and a target Q = dG, and return (d, stats).
"""
from __future__ import annotations
import math, random
from .ec import Curve


def bsgs(E: Curve, G, Q, n: int | None = None, m: int | None = None):
    """Baby-step giant-step. Memory m (default ⌈√n⌉) entries; time ≈ m + n/m group operations."""
    n = n or E.n
    m = m or math.isqrt(n) + 1
    E.ops.reset()
    table = {}
    P = None
    for j in range(m):               # baby steps jG, j = 0..m-1
        key = P
        table.setdefault(key, j)
        P = E.add(P, G)
    mG_neg = E.neg(P)                # −mG
    R = Q
    for i in range(0, n // m + 2):   # giant steps Q − i·mG
        if R in table:
            d = (i * m + table[R]) % n
            return d, {"ops": E.ops.as_dict(), "table_entries": len(table), "giant_steps": i + 1, "memory_entries": m}
        R = E.add(R, mG_neg)
    raise ValueError("BSGS failed (Q not in <G>?)")


class RAddingWalk:
    """Teske's r-adding walk f(P) = P + M_{h(P)}, M_j = u_j G + v_j Q, with scalar bookkeeping (a, b): P = aG + bQ."""

    def __init__(self, E: Curve, G, Q, n: int, r: int, rng: random.Random):
        self.E, self.G, self.Q, self.n, self.r = E, G, Q, n, r
        self.u = [rng.randrange(1, n) for _ in range(r)]
        self.v = [rng.randrange(1, n) for _ in range(r)]
        self.M = [E.add(E.mul(u, G), E.mul(v, Q)) for u, v in zip(self.u, self.v)]

    def partition(self, P) -> int:
        # Hash of the x-coordinate; a multiplicative hash avoids correlation with the low bits of x.
        return (P[0] * 0x9E3779B97F4A7C15 >> 7) % self.r

    def step(self, P, a, b):
        j = self.partition(P)
        return self.E.add(P, self.M[j]), (a + self.u[j]) % self.n, (b + self.v[j]) % self.n


def pollard_rho(E: Curve, G, Q, n: int | None = None, r: int = 20, seed: int = 0, max_iter: int | None = None):
    """Single-walk Pollard rho with Floyd cycle detection (reference). Returns (d, stats)."""
    n = n or E.n
    rng = random.Random(seed)
    E.ops.reset()
    walk = RAddingWalk(E, G, Q, n, r, rng)
    a0, b0 = rng.randrange(n), rng.randrange(n)
    X = E.add(E.mul(a0, G), E.mul(b0, Q))
    if X is None:
        return (-a0 * pow(b0, -1, n)) % n, {"ops": E.ops.as_dict(), "iterations": 0}
    E.ops.reset()
    tort = (X, a0, b0)
    hare = (X, a0, b0)
    it = 0
    while True:
        tort = walk.step(*tort)
        hare = walk.step(*walk.step(*hare))
        it += 1
        if max_iter and it > max_iter:
            raise RuntimeError("rho exceeded max_iter")
        if tort[0] is None or hare[0] is None:
            # landed on infinity: restart with new seed
            return pollard_rho(E, G, Q, n, r, seed + 1_000_003, max_iter)
        if tort[0] == hare[0]:
            a1, b1 = tort[1], tort[2]
            a2, b2 = hare[1], hare[2]
            if (b1 - b2) % n == 0:
                return pollard_rho(E, G, Q, n, r, seed + 1_000_003, max_iter)  # useless collision
            d = (a2 - a1) * pow(b1 - b2, -1, n) % n
            assert E.mul(d, G) == Q
            return d, {"ops": E.ops.as_dict(), "iterations": it, "floyd_steps": 3 * it}


def pollard_rho_dp(E: Curve, G, Q, n: int | None = None, r: int = 20, dp_bits: int = 4, walks: int = 4,
                   seed: int = 0, neg_map: bool = False, max_steps: int | None = None):
    """Distinguished-point (van Oorschot–Wiener) rho with `walks` sequential pseudo-parallel walks.

    If neg_map is True the walk acts on classes {P, −P} (canonical representative: smaller y) and the usual
    2-cycle look-ahead escape is used; fruitless cycles longer than 2 are broken by a bounded-length check.
    """
    n = n or E.n
    rng = random.Random(seed)
    E.ops.reset()
    walk = RAddingWalk(E, G, Q, n, r, rng)
    mask = (1 << dp_bits) - 1
    table = {}
    stats = {"fruitless_2cycles": 0, "fruitless_long": 0, "dp_count": 0, "iterations": 0, "restarts": 0}

    def canon(P, a, b):
        if P[1] > E.p - P[1]:
            return (P[0], E.p - P[1]), (-a) % n, (-b) % n
        return P, a, b

    def new_walk():
        a, b = rng.randrange(n), rng.randrange(n)
        P = E.add(E.mul(a, G), E.mul(b, Q))
        while P is None:
            a, b = rng.randrange(n), rng.randrange(n)
            P = E.add(E.mul(a, G), E.mul(b, Q))
        return canon(P, a, b) if neg_map else (P, a, b)

    state = [new_walk() for _ in range(walks)]
    steps_since_check = [0] * walks
    anchor = [None] * walks
    while True:
        for w in range(walks):
            P, a, b = state[w]
            if neg_map:
                j = walk.partition(P)
                P2 = E.add(P, walk.M[j])
                if P2 is None:
                    state[w] = new_walk(); stats["restarts"] += 1; continue
                P2, a2, b2 = canon(P2, (a + walk.u[j]) % n, (b + walk.v[j]) % n)
                # 2-cycle look-ahead (Bos–Kleinjung–Lenstra 2010): if f(P2) == P, use the next partition
                j2 = walk.partition(P2)
                P3 = E.add(P2, walk.M[j2])
                if P3 is not None and canon(P3, 0, 0)[0] == P:
                    stats["fruitless_2cycles"] += 1
                    j = (j + 1) % r
                    P2 = E.add(P, walk.M[j])
                    if P2 is None:
                        state[w] = new_walk(); stats["restarts"] += 1; continue
                    P2, a2, b2 = canon(P2, (a + walk.u[j]) % n, (b + walk.v[j]) % n)
                P, a, b = P2, a2, b2
                # long fruitless cycles: compare with an anchor point every 4096 steps
                steps_since_check[w] += 1
                if steps_since_check[w] % 4096 == 0:
                    anchor[w] = P
                elif anchor[w] is not None and P == anchor[w]:
                    stats["fruitless_long"] += 1
                    # escape by doubling
                    P = E.double(P); a, b = 2 * a % n, 2 * b % n
                    P, a, b = canon(P, a, b)
                    anchor[w] = None
            else:
                P, a, b = walk.step(P, a, b)
                if P is None:
                    state[w] = new_walk(); stats["restarts"] += 1; continue
            stats["iterations"] += 1
            if max_steps and stats["iterations"] > max_steps:
                raise RuntimeError("exceeded max_steps")
            if P[0] & mask == 0:
                stats["dp_count"] += 1
                if P in table:
                    a1, b1 = table[P]
                    if (b1 - b) % n != 0:
                        d = (a - a1) * pow(b1 - b, -1, n) % n
                        assert E.mul(d, G) == Q
                        stats["ops"] = E.ops.as_dict()
                        stats["table_entries"] = len(table)
                        return d, stats
                    # same walk merged with itself (useless); restart this walk
                    state[w] = new_walk(); stats["restarts"] += 1; continue
                table[P] = (a, b)
                state[w] = new_walk()  # standard vOW: after a DP, start a fresh walk
                continue
            state[w] = (P, a, b)


def kangaroo(E: Curve, G, Q, lo: int, hi: int, n: int | None = None, seed: int = 0):
    """Pollard's kangaroo (lambda) method for d ∈ [lo, hi]: one tame and one wild kangaroo, DP-free
    'van Oorschot–Wiener trap' variant (tame kangaroo lays a trap, wild follows). Expected ≈ 2√w operations."""
    n = n or E.n
    w = hi - lo + 1
    rng = random.Random(seed)
    E.ops.reset()
    k = max(1, int(math.log2(math.sqrt(w))))  # jumps 2^0 .. 2^k, mean ≈ √w/2 (Pollard 2000)
    jumps = [1 << i for i in range(k + 1)]
    J = [E.mul(j, G) for j in jumps]
    mean = sum(jumps) / len(jumps)
    N = int(2 * math.sqrt(w) / mean * 4) + 10  # tame path length (van Oorschot–Wiener: ~2√w total)

    salt = rng.getrandbits(64) | 1

    def h(P):
        if P is None:  # the walk may pass through the point at infinity when the interval is the whole group
            return 0
        return ((P[0] ^ salt) * 0x9E3779B97F4A7C15 >> 11) % len(jumps)

    # tame kangaroo starts at the upper end of the interval and lays a trap at every point it visits
    t = hi
    T = E.mul(t, G)
    traps = {}
    for _ in range(N):
        traps[T] = t
        j = h(T)
        T = E.add(T, J[j]); t += jumps[j]
    traps[T] = t
    tame_dist = t - hi
    # wild kangaroo starts at Q = dG; it is caught once it lands on a trap. It passes the whole tame path
    # after at most w + tame_dist distance, so beyond that the attempt has failed.
    Wp, wdist = Q, 0
    it = 0
    while wdist <= w + tame_dist:
        if Wp in traps:
            d = (traps[Wp] - wdist) % n
            if E.mul(d, G) == Q:
                return d, {"ops": E.ops.as_dict(), "iterations": N + it, "tame_steps": N, "wild_steps": it}
        j = h(Wp)
        Wp = E.add(Wp, J[j]); wdist += jumps[j]
        it += 1
    # failure: retry with a different seed/longer tame path
    return kangaroo(E, G, Q, lo, hi, n, seed + 1)


def pohlig_hellman(E: Curve, G, Q, N: int, seed: int = 0):
    """Pohlig–Hellman for a point G of composite order N (smooth N makes this cheap); uses BSGS in each prime-power subgroup."""
    from sympy import factorint
    from sympy.ntheory.modular import crt
    rems, mods = [], []
    for q, e in factorint(N).items():
        qe = q ** e
        Gq = E.mul(N // qe, G)
        Qq = E.mul(N // qe, Q)
        # solve in the subgroup of order q^e digit by digit
        x = 0
        Gq1 = E.mul(q ** (e - 1), Gq)  # order q
        for k in range(e):
            Qk = E.mul(q ** (e - 1 - k), E.sub(Qq, E.mul(x, Gq)))
            dk, _ = bsgs(E, Gq1, Qk, n=q)
            x += dk * q ** k
        rems.append(x); mods.append(qe)
    d, _ = crt(mods, rems)
    return int(d) % N, {}