#!/usr/bin/env python
# coding: utf-8

# This is a code for point counting on genus 3 curves of type y^2 = x^7 + a*x^4 + bx over prime field Fp.
#
# Tested on SageMath v8.9 with Python 2.

import time

# Auxiliary functions
R1.<T> = ZZ['T']
R2.<a1,a2,a3> = ZZ["a1,a2,a3"]

def EC(f):
    l_a2 = f.coefficients(sparse=False)[2]
    l_a4 = f.coefficients(sparse=False)[1]
    l_a6 = f.coefficients(sparse=False)[0]
    try:
        res = EllipticCurve([0,l_a2,0,l_a4,l_a6])
    except ArithmeticError:
        return 'singular'
    return res

def HEC(f):
    try:
        res = HyperellipticCurve(f)
    except ValueError:
        return 'singular'
    return res

def is_nth_root(number, N):
    try:
        number.nth_root(N)
        return True
    except ValueError:
        return False

# Generation of a random point on a hyperelliptic curve C
def HEC_random_point(C):
    f = C.hyperelliptic_polynomials()[0]
    while True:
        x_r = f.base_ring().random_element()
        y2 = f(x_r)
        if y2.is_square():
            return [x_r, y2.sqrt()]

# Generation of n random unique points on a hyperelliptic curve C
def HEC_random_points_uniq(C, n):
    f = C.hyperelliptic_polynomials()[0]
    res = []
    for i in range(1,n+1):
        tries = 100
        found = False
        while tries > 0 and (not found):
            P = HEC_random_point(C)
            if not (P in res):
                res.append(P)
                found = True
            tries = tries - 1
    return res

# Generation of a random element in Jacobian of a curve C
def JC_random_element(C):
    f = C.hyperelliptic_polynomials()[0]
    R.<x> = f.base_ring()['x']
    J = C.jacobian()
    points = HEC_random_points_uniq(C, C.genus())
    u = 1
    for point in points:
        u = u * (x-point[0])
    v = f.parent().lagrange_polynomial(points)
    return J([u, v])


# Testing of the Jacobian order candidates
def check_order(C, N):
    f = C.hyperelliptic_polynomials()[0]
    q = f.base_ring().cardinality()
    g = C.genus()
    # Check Hasse-Weil bound
    if not ( (q.sqrt() - 1)^(2*g) <= N and N <= (q.sqrt() + 1)^(2*g) ):
        return False
    res = True
    # Check by multiplying on random elements of the Jacobian
    i = 1
    while i<10 and res:
        P = JC_random_element(C)
        if list(P*N)!=list(C.jacobian()(0)):
            res = False
            break
        i = i + 1
    return res

# Computation of the characteristic polynomial chi_{C,p}(T) of a curve C: y^2 = x^7 + a*x^4 + b x
def g3_frobenius_polynomial(a, b):
    Fq = a.parent()
    q = Fq.cardinality()
    p = Fq.characteristic()
    #Fq2 = Fq.extension(2)
    R.<x> = Fq['x']
    f = x^7 + a*x^4 + b*x
    C = HEC(f)
    #print("-> p (mod 3) =", Mod(p, 3))
    #print("-> b is a square root:", is_nth_root(b, 2))
    #print("-> b is a 3-root:", is_nth_root(b, 3))
    #print("-> -1 is a square:", is_nth_root(Fq(-1), 2))
    
    frobs = []
    E1 = EC(x^3 + a*x^2 + b*x)
    E2t = EC(x^3 - 3*b*x + a*b)

    tmc = time.time()
    print("Computing elliptic curves traces ... ")
    t1 = q + 1 - E1.cardinality()
    tt2 = q + 1 - E2t.cardinality()
    tmc = time.time()-tmc
    print("-> t1 ={}".format(t1))
    print("-> tt2 = {}".format(tt2))
    if Mod(t1,p) == 0:
        print("-> E1 is supersingualar")
    else:
        print("-> E1 is ordinary")
    if Mod(tt2,p) == 0:
        print("-> E2/E2t is supersingualar")
    else:
        print("-> E2/E2t is ordinary")
    print("done in {} sec.".format(tmc))

    print("Checking possible frobs ... ")
    tmc = time.time()
    candidates = [
        (T^2 - t1*T + q) * (T^2 - tt2*T + q) * (T^2 + tt2*T + q),
        (T^2 - t1*T + q) * (T^2 - tt2*T + q) * (T^2 - tt2*T + q),
        (T^2 - t1*T + q) * (T^4 + tt2*T^3 + (tt2^2 - q)*T^2 + tt2*q*T + q^2),
        (T^2 - t1*T + q) * (T^4 - tt2*T^3 + (tt2^2 - q)*T^2 - tt2*q*T + q^2)
    ]
    frobs = []
    for cand in candidates:
        N = cand(1)
        if check_order(C, N):
            frobs = frobs + [cand]
    tmc = time.time()-tmc 
    print("-> len of list: {}".format(len(frobs)))
    print("done in {} sec.".format(tmc))
    #print "[debug] real polynomial: ", C.frobenius_polynomial()
    return list(set(frobs))


if len(sys.argv) < 4:
    print("This is a script for point-counting on the curve y^2 = x^7 + a*x^4 + bx over prime field F_p.")
    print("Usage:\n\tsage {} p a b".format(sys.argv[0][0:-3]))
    exit()

p = ZZ(sys.argv[1])
q = p
Fq = FiniteField(q)
R.<x> = Fq['x']
ZZx.<T> = ZZ['T']

a = Fq(sys.argv[2])
b = Fq(sys.argv[3])

print("p = {}".format(p))
print("p (mod 3) = {}".format(p % 3))
print("a = {}".format(a))
print("b = {}".format(b))
#print("b is a cubic residue: {}".format(is_nth_root(b,3)))
print("memory usage (begin): " + str(get_memory_usage()))
tm = time.time()
frobs = g3_frobenius_polynomial(a,b)
tm = time.time() - tm
print("Frobenius polynomial: {}".format(frobs[0]))
print("-> computed in {} sec.".format(tm))
print("-> memory usage (at this point): " + str(get_memory_usage()))
frob_fact = frobs[0].factor()
print("Frobenius polynomial (fact.): {}".format(frob_fact))
print("List size: {}".format(len(frobs)))
#print("Full List: {}".format(frobs))
#print("-> factorization: {}".format(factor(N, proof=False)))
if len(frob_fact) == 2 and frob_fact[1][1] == 1:
    frob_A = frob_fact[1][0]
    print("frob_A = {}".format(frob_A))
    N_A = frob_A(1)
    print("#A(Fq) = {} ({} bit)".format(N_A, floor(log(N_A)/log(2))))
    print("-> is_pseudoprime(#A)? {}".format(is_pseudoprime(N_A)))
    #print("-> factorization: {}".format(factor(N_A)))
N = frobs[0](1)
print("#J(Fq) = {} ({} bit)".format(N, floor(log(N)/log(2))))