#!/usr/bin/env python3
import json
from math import log2
from random import randint
from time import perf_counter
from collections import defaultdict
from contextlib import contextmanager
from sage.all import random_prime, is_prime, GF, EllipticCurve, sqrt as sage_sqrt, primitive_root

STATS = defaultdict(lambda: defaultdict(int))
COUNT = 25
BITLENS = range(30,150)

@contextmanager
def measure(key, action):
    start = perf_counter()
    try:
        yield
    finally:
        STATS[key][action] += perf_counter() - start

def check_glv_endomorphism(curve, generator=None):
    """
    Check if an elliptic curve supports the specific endomorphism needed for GLV method.

    Parameters:
        curve: An elliptic curve defined over a finite field
        generator: Optional generator point. If None, will use the default generator

    Returns:
        A dictionary containing the test results and relevant values
    """
    # Get the curve parameters
    F = curve.base_field()
    p = F.order()
    n = curve.order()

    # Check if curve has the correct form for j-invariant 0
    if curve.a_invariants() != (0, 0, 0, 0, curve.a_invariants()[4]):
        return {"supports_glv": False, "reason": "Curve is not in the form y^2 = x^3 + b"}

    # Check field characteristics
    if p % 3 != 1:
        return {"supports_glv": False, "reason": "Field characteristic p must satisfy p ≡ 1 (mod 3)"}

    # Check if n is prime
    if not n.is_prime():
        return {"supports_glv": False, "reason": "Curve order must be prime"}

    # Use provided generator or get default
    if generator is None:
        try:
            generator = curve.gens()[0]
        except:
            return {"supports_glv": False, "reason": "No generator specified and unable to find default generator"}

    # Find a cube root of unity in the base field
    try:
        # primitive_root() is usually faster than via F.unit_group(), e.g. F_star = F.unit_group(), g_base = F(F_star.gen(0))
        g_base = F(primitive_root(F.order()))
        beta = g_base**((p-1)//3)

        # Verify beta is a non-trivial cube root of unity
        if beta == 1 or beta**3 != 1:
            return {"supports_glv": False, "reason": "Failed to find valid cube root of unity in base field"}

        # Verify beta satisfies the minimal polynomial
        if beta**2 + beta + 1 != 0:
            return {"supports_glv": False, "reason": "Cube root of unity does not satisfy minimal polynomial"}
    except Exception as e:
        return {"supports_glv": False, "reason": f"Error finding cube root of unity: {e}"}

    # Find primitive 6th root of unity in the scalar field
    try:
        Fq = GF(n)
        g_scalar = Fq(primitive_root(n))

        # Try both (n-1)/6 and (n-1)/3 for lambda
        lambda_val_1 = g_scalar**((n-1)//6)
        lambda_val_2 = g_scalar**((n-1)//3)

        # Test both values to see if either works
        lambda_vals = [lambda_val_1, lambda_val_2]
        working_lambda = None

        for lambda_val in lambda_vals:
            # Verify lambda^6 = 1
            if lambda_val**6 != 1:
                continue

            # Check if lambda satisfies the minimal polynomial
            if lambda_val**2 + lambda_val + 1 != 0:
                continue

            # The critical test: check if lambda*G = (beta*G_x, G_y)
            try:
                endo_point = curve(beta * generator[0], generator[1])
                scalar_point = lambda_val * generator

                if endo_point == scalar_point:
                    working_lambda = lambda_val
                    break
            except Exception:
                continue

        if working_lambda is None:
            return {
                "supports_glv": False,
                "reason": "No suitable lambda value found that satisfies the endomorphism relation",
                "beta": beta,
                "lambda_candidates": lambda_vals,
                "test_points": {
                    "endomorphism": endo_point,
                    "scalar_mult": scalar_point
                }
            }

    except Exception as e:
        return {"supports_glv": False, "reason": f"Error finding lambda value: {e}"}

    # If we get here, the curve supports GLV
    return {
        "supports_glv": True,
        "beta": beta,
        "lambda": working_lambda,
        "verification": "lambda*G == E(beta*G[0], G[1])"
    }

def find_generator(E:EllipticCurve):
    x = E.base_field()(1)
    while True:
        try:
            G = E.lift_x(x)
        except ValueError:
            x += 1
            continue
        if G.order() == E.order():
            # Generator must be even?
            if int(G[1]) & 1:
                G = -G
            return G
        x += 1

def find_split_constants_explicit_tof(p:int, E:EllipticCurve):
    # from libsecp256k1/sage
    assert p % 3 == 1
    assert E.j_invariant() == 0
    t = int(E.trace_of_frobenius())
    c = int(sage_sqrt((4*p - t**2)//3))
    b1, b2 = c, (1 - (t - c)//2) % E.order()
    return b1, int(b2)

def decompose_scalar(k, n, g1, g2, b1, b2, lambda_val, bitsize2=384):
    """Decompose scalar k into k1 and k2 using GLV method."""
    k = k % n

    # Compute c1 and c2 using precomputed constants and bit operations
    # This replaces division with multiplication and right shift
    c1 = (k * g1) >> bitsize2
    c2 = (k * g2) >> bitsize2

    # Handle final bit for rounding as per the paper
    if (k * g1) & (1 << (bitsize2-1)):
        c1 += 1
    if (k * g2) & (1 << (bitsize2-1)):
        c2 += 1

    neg_b1 = -b1 % n
    neg_b2 = -b2 % n

    # Compute k2 = -c1·(-b1) - c2·(-b2)
    k2 = (c1 * neg_b1 + c2 * neg_b2) % n

    # Compute k1 = k - k2·λ mod n
    k1 = (k - (k2 * lambda_val) % n) % n

    # Ensure k1 and k2 are properly minimized
    if k1 > (n >> 1):
        k1 = k1 - n
    if k2 > (n >> 1):
        k2 = k2 - n

    return k1, k2

def find_g1_g2(b1, b2, n, bitsize=256, bitsize2=384):
    # Python's round() is off by 1
    quotient_g1 = (2**bitsize2)*(-b2)//n
    remainder_g1 = (2**bitsize2)*(-b2)%n
    g1 = (quotient_g1 + (1 if remainder_g1 >= n//2 else 0)) % (2**bitsize)
    # Python's round() is off by one
    quotient_g2 = (2**bitsize2)*(b1)//n
    remainder_g2 = (2**bitsize2)*(b1)%n
    g2 = (quotient_g2 + (1 if remainder_g2 >= n//2 else 0)) % (2**bitsize)
    return g1, g2

def sample_primes(bitlen):
    while True:
        with measure(bitlen, '1_time_random_prime'):
            p = random_prime(2**bitlen, lbound=2**(bitlen-1))
        if p % 12 == 7:
            yield p
            STATS[bitlen]['1_prime_congruent'] += 1
            continue
        STATS[bitlen]['1_prime_not_congruent'] += 1

def sample_curves(bitlen):
    for p in sample_primes(bitlen):
        with measure(bitlen, '2_time_primitive_root'):
            F = GF(p)
            g = F(primitive_root(p))
        n_found = 0
        # 6th primitive roots g^1 and g^-1 are only ones which can be prime
        for b in [int(F(g)), int(F(g**5))]:
            with measure(bitlen, '2_time_check_curve'):
                E = EllipticCurve(F, [0, b])
                E_q = E.order()
                E_q_prime = is_prime(E_q)
            if E_q_prime:
                n_found += 1
                yield p, b, E
        if n_found == 0:
            STATS[bitlen]['2_no_prime_curve'] += 1
        else:
            STATS[bitlen]['2_prime_curve'] += 1
            if n_found == 2:
                STATS[bitlen]['2_1_prime_curve_both'] += 1
            else:
                STATS[bitlen]['2_1_prime_curve_one'] += 1

def try_decompose(E_q, bitlen, b1, b2, lambda_val):
    g1, g2 = find_g1_g2(b1, b2, E_q, bitlen, bitlen//2 + bitlen)
    count = 10
    k1_log2_total = 0
    k2_log2_total = 0
    for _ in range(count):
        k = randint(2**510, 2**512) % E_q
        k1, k2 = decompose_scalar(k, E_q, g1, g2, b1, b2, lambda_val, (bitlen//2)+bitlen)
        recomposed = (k1 + (k2 * lambda_val) % E_q) % E_q
        if k != recomposed:
            return None
        k1_log2 = log2(abs(int(k1))) if k1 != 0 else 0
        k2_log2 = log2(abs(int(k2))) if k2 != 0 else 0
        k1_log2_total += k1_log2
        k2_log2_total += k2_log2
    k1_log2 = round(k1_log2_total / count,2)
    k2_log2 = round(k2_log2_total / count,2)
    return k1_log2, k2_log2, k1_log2 <= bitlen//2 and k2_log2 <= bitlen//2, int(lambda_val) & 1

with open('graphs/3-glv-decomp.jsonl', 'w') as handle:
    for bitlen in BITLENS:
        FOUND = 0
        bitlen_start = perf_counter()
        with measure(bitlen, 'time_total'):
            for p, b, E in sample_curves(bitlen):
                if FOUND >= COUNT:
                    break
                with measure(bitlen, '2_time_basic_field_checks'):
                    b_log2 = log2(b)
                    E_q = E.order()
                    E_q_prime = is_prime(E_q)

                STATS[bitlen]['3_found_curves'] += 1

                with measure(bitlen, '3_time_find_generator'):
                    G = find_generator(E)
                if log2(G[1]) <= 8:
                    STATS[bitlen]['3_1_G_y_is_under_8bit'] += 1
                else:
                    STATS[bitlen]['3_1_G_y_is_over_8bit'] += 1
                with measure(bitlen, '3_time_check_glv_endomorphism'):
                    glv_info = check_glv_endomorphism(E, G)
                supports_glv = glv_info['supports_glv']
                if not supports_glv:
                    STATS[bitlen]['3_not_glv'] += 1
                    continue
                STATS[bitlen]['3_glv'] += 1

                lambda_val = glv_info['lambda']
                neg_lambda = -int(lambda_val) % E_q

                with measure(bitlen, '4_time_find_split_constants_explicit_tof'):
                    b1, b2 = find_split_constants_explicit_tof(p, E)

                ok = False
                ok_i = None
                with measure(bitlen, '4_time_try_decompose'):
                    checks = [
                        try_decompose(E_q, bitlen, b1, b2, lambda_val),
                        try_decompose(E_q, bitlen, -b1 % E_q, b2, lambda_val),
                        try_decompose(E_q, bitlen, b1, -b2 % E_q, lambda_val),
                        try_decompose(E_q, bitlen, -b1 % E_q, -b2 % E_q, lambda_val),
                        try_decompose(E_q, bitlen, b1, b2, neg_lambda),
                        try_decompose(E_q, bitlen, -b1 % E_q, b2, neg_lambda),
                        try_decompose(E_q, bitlen, b1, -b2 % E_q, neg_lambda),
                        try_decompose(E_q, bitlen, -b1 % E_q, -b2 % E_q, neg_lambda)
                    ]
                for i, x in enumerate(checks):
                    if x[2]:
                        ok = True
                        ok_i = i

                if ok:
                    STATS[bitlen]['4_glv_efficient_decompose'] += 1
                    print()
                    print(f"{bitlen}-bit Curve with b = {b} (log2={round(b_log2,3)},{round(b_log2/log2(p),3)}), j-invariant = {E.j_invariant()}, order = {E_q} (prime? {E_q_prime})")
                    print(' - G', G)
                    print(' - GLV', glv_info)
                    for i, x in enumerate(checks):
                        print(i, x)
                    print(ok)
                    print()
                    print()
                    FOUND += 1
                else:
                    STATS[bitlen]['4_glv_inefficient_decompose'] += 1

        for b, x in STATS.items():
            for y in sorted(x.keys()):
                print(b, y, x[y])
        print()

        handle.write(json.dumps((bitlen,STATS[bitlen])) + "\n")
        handle.flush()