"""Short Weierstrass elliptic curves over prime fields with exact integer arithmetic.

Points are affine tuples (x, y) or None for the point at infinity. Every group operation is counted in an
OpCounter so that experiments can report operation counts alongside wall-clock time.
"""
from __future__ import annotations
import random
from dataclasses import dataclass, field


class OpCounter:
    __slots__ = ("add", "dbl", "inv", "mul")

    def __init__(self):
        self.add = self.dbl = self.inv = self.mul = 0

    def reset(self):
        self.add = self.dbl = self.inv = self.mul = 0

    def group_ops(self) -> int:
        return self.add + self.dbl

    def as_dict(self) -> dict:
        return {"add": self.add, "dbl": self.dbl, "inv": self.inv, "mul": self.mul}


def legendre(a: int, p: int) -> int:
    return pow(a % p, (p - 1) // 2, p)


def sqrt_mod(a: int, p: int):
    """Tonelli–Shanks square root modulo an odd prime p, or None if a is a non-residue."""
    a %= p
    if a == 0:
        return 0
    if legendre(a, p) != 1:
        return None
    if p % 4 == 3:
        return pow(a, (p + 1) // 4, p)
    q, s = p - 1, 0
    while q % 2 == 0:
        q //= 2
        s += 1
    z = 2
    while legendre(z, p) != p - 1:
        z += 1
    m, c, t, r = s, pow(z, q, p), pow(a, q, p), pow(a, (q + 1) // 2, p)
    while t != 1:
        i, t2 = 0, t
        while t2 != 1:
            t2 = t2 * t2 % p
            i += 1
        b = pow(c, 1 << (m - i - 1), p)
        m, c, t, r = i, b * b % p, t * b * b % p, r * b % p
    return r


@dataclass
class Curve:
    p: int
    a: int
    b: int
    n: int | None = None          # order of the subgroup generated by G (prime in our instances)
    h: int = 1                    # cofactor
    G: tuple | None = None
    name: str = ""
    ops: OpCounter = field(default_factory=OpCounter, repr=False, compare=False)

    # ---- curve invariants -------------------------------------------------------------------
    def discriminant(self) -> int:
        return (-16 * (4 * pow(self.a, 3, self.p) + 27 * pow(self.b, 2, self.p))) % self.p

    def j_invariant(self) -> int:
        a3 = 4 * pow(self.a, 3, self.p)
        den = (a3 + 27 * pow(self.b, 2, self.p)) % self.p
        return 1728 * a3 * pow(den, -1, self.p) % self.p  # j = 1728 · 4a³ / (4a³ + 27b²)

    def is_on_curve(self, P) -> bool:
        if P is None:
            return True
        x, y = P
        return (y * y - (x * x * x + self.a * x + self.b)) % self.p == 0

    # ---- group law ------------------------------------------------------------------------------
    def neg(self, P):
        return None if P is None else (P[0], (-P[1]) % self.p)

    def add(self, P, Q):
        if P is None:
            return Q
        if Q is None:
            return P
        p = self.p
        x1, y1 = P
        x2, y2 = Q
        if x1 == x2:
            if (y1 + y2) % p == 0:
                return None
            return self.double(P)
        self.ops.add += 1
        self.ops.inv += 1
        lam = (y2 - y1) * pow(x2 - x1, -1, p) % p
        x3 = (lam * lam - x1 - x2) % p
        y3 = (lam * (x1 - x3) - y1) % p
        return (x3, y3)

    def double(self, P):
        if P is None:
            return None
        p = self.p
        x1, y1 = P
        if y1 == 0:
            return None
        self.ops.dbl += 1
        self.ops.inv += 1
        lam = (3 * x1 * x1 + self.a) * pow(2 * y1, -1, p) % p
        x3 = (lam * lam - 2 * x1) % p
        y3 = (lam * (x1 - x3) - y1) % p
        return (x3, y3)

    def sub(self, P, Q):
        return self.add(P, self.neg(Q))

    def mul(self, k: int, P):
        """Left-to-right double-and-add. k may be any integer (negative allowed)."""
        if P is None or k == 0:
            return None
        if k < 0:
            k, P = -k, self.neg(P)
        R = None
        for bit in bin(k)[2:]:
            R = self.double(R)
            if bit == "1":
                R = self.add(R, P)
        return R

    def mul_wnaf(self, k: int, P, w: int = 4):
        """Signed-window (wNAF) scalar multiplication; same result as mul, fewer additions."""
        if P is None or k == 0:
            return None
        if k < 0:
            k, P = -k, self.neg(P)
        # precompute odd multiples 1P, 3P, ..., (2^{w-1}-1)P
        twoP = self.double(P)
        table = [P]
        for _ in range((1 << (w - 2)) - 1):
            table.append(self.add(table[-1], twoP))
        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
        R = None
        for d in reversed(naf):
            R = self.double(R)
            if d > 0:
                R = self.add(R, table[(d - 1) // 2])
            elif d < 0:
                R = self.add(R, self.neg(table[(-d - 1) // 2]))
        return R

    # ---- utilities ----------------------------------------------------------------------------
    def random_point(self, rng: random.Random):
        while True:
            x = rng.randrange(self.p)
            rhs = (x * x * x + self.a * x + self.b) % self.p
            y = sqrt_mod(rhs, self.p)
            if y is not None:
                return (x, y if rng.random() < 0.5 else (-y) % self.p)

    def lift_x(self, x: int, even_y: bool = True):
        rhs = (x * x * x + self.a * x + self.b) % self.p
        y = sqrt_mod(rhs, self.p)
        if y is None:
            return None
        if (y % 2 == 0) != even_y:
            y = (-y) % self.p
        return (x, y)

    def order_of_point(self, P, bound: int | None = None) -> int:
        """Naive order (only for tiny curves); uses n if known."""
        if self.n is not None and self.mul(self.n, P) is None:
            # smallest divisor of n·h that kills P
            N = self.n * self.h
            for d in sorted(_divisors(N)):
                if self.mul(d, P) is None:
                    return d
        k, R = 1, P
        while R is not None:
            R = self.add(R, P)
            k += 1
            if bound and k > bound:
                raise ValueError("order bound exceeded")
        return k

    def count_points_naive(self) -> int:
        """#E(F_p) by Legendre symbols — only for p < ~10^6."""
        total = 1  # infinity
        for x in range(self.p):
            rhs = (x * x * x + self.a * x + self.b) % self.p
            if rhs == 0:
                total += 1
            elif legendre(rhs, self.p) == 1:
                total += 2
        return total


def _divisors(N: int):
    from sympy import divisors
    return divisors(N)


# ---- synthetic curve generation via PARI/GP (SEA point counting) ---------------------------------------
def gp_ellcard_batch(curves: list[tuple[int, int, int]]) -> list[int]:
    """Orders #E(F_p) for a list of (p, a, b) using PARI/GP's ellcard (SEA)."""
    import subprocess
    lines = ["default(parisize, 64000000);"]
    for (p, a, b) in curves:
        lines.append(f"print(ellcard(ellinit([{a},{b}],{p})));")
    out = subprocess.run(["gp", "-q", "-f"], input="\n".join(lines) + "\n", capture_output=True, text=True,
                         timeout=600)
    if out.returncode != 0:
        raise RuntimeError(out.stderr)
    return [int(s) for s in out.stdout.split()]


def random_prime_order_curve(bits: int, rng: random.Random, j0: bool = False, batch: int = 32) -> Curve:
    """Random curve y² = x³ + ax + b over a random `bits`-bit prime p with prime group order.

    If j0 is True, a = 0 and p ≡ 1 (mod 3) so that the curve has the order-6 automorphism group of secp256k1.
    """
    from sympy import isprime, randprime
    while True:
        if j0:
            while True:
                p = randprime(1 << (bits - 1), 1 << bits)
                if p % 3 == 1:
                    break
        else:
            p = randprime(1 << (bits - 1), 1 << bits)
        cands = []
        for _ in range(batch):
            a = 0 if j0 else rng.randrange(p)
            b = rng.randrange(1, p)
            if (4 * a ** 3 + 27 * b ** 2) % p == 0:
                continue
            cands.append((p, a, b))
        orders = gp_ellcard_batch(cands)
        for (p, a, b), N in zip(cands, orders):
            if isprime(N) and N != p:
                E = Curve(p, a, b, n=N, h=1, name=f"synthetic-{bits}b{'-j0' if j0 else ''}")
                E.G = E.random_point(rng)
                assert E.mul(N, E.G) is None
                return E


def cube_roots_of_unity(m: int, rng: random.Random | None = None) -> list[int]:
    """The two primitive cube roots of unity modulo a prime m ≡ 1 (mod 3)."""
    assert m % 3 == 1
    rng = rng or random.Random(0)
    while True:
        g = rng.randrange(2, m)
        beta = pow(g, (m - 1) // 3, m)
        if beta != 1:
            return sorted([beta, beta * beta % m])


def endomorphism_lambda(E: Curve, beta: int) -> int:
    """Return λ with (βx, y) = [λ]P on E (prime order n ≡ 1 mod 3, a = 0)."""
    assert E.a == 0 and E.n is not None and E.n % 3 == 1
    lams = cube_roots_of_unity(E.n)
    P = E.G
    phiP = (beta * P[0] % E.p, P[1])
    for lam in lams:
        if E.mul(lam, P) == phiP:
            return lam
    raise ValueError("no λ matches β")