# Number of trials to perform:

NTRIES = 10

# Case N=4, uncomment to test Conjecture 1
#A = 2
#B = 3
#C = 6
#K = 1
#N = 4

# Case N=5, uncomment to test Conjectures 2,3
A = 2
B = 3
C = 6
K = 2
N = 5

# CF is the field Q[zeta], zeta is the primitive
# C's root of unity
CF.<zeta> = CyclotomicField(C)

# Uncomment to test Conjecture 1:
#DIVISOR = (0,1,2,1)
# Uncomment to test Conjecture 2:
# DIVISOR = None
# Uncomment to test Conjecture 3:
DIVISOR = (1,0,zeta,0,-1)

#DIVISOR = (zeta, zeta - 1, -zeta, 0, 0)


A1 = K*C/A
B1 = K*C/B
A2 = C/A
B2 = C/B
assert C==lcm(A, B)
G = gcd(A, B)
D = K*C^2/A/B

# Recall the setup from https://arxiv.org/abs/1705.02931:
# We study the curve x^A-y^B=z^C, which lives in P(A2,B2,1)
# In affine coordinates, it is given by x^A-y^B=1
# A1, B1 are the degrees of the polynomials P_i, Q_i
# which are be used to define a rational curve in E^N/\Gamma


# If FINITE_FIELD is False, we perform computations over Q where
# the coordinates of a test point are variables (too slow and
# doesn't work in practice, but implemented for convenience)
# If FINITE_FIELD is True, perform computations over a finite field
# for a random point.
FINITE_FIELD = True

if FINITE_FIELD:
    # Our favorite finite field containing C's roots of unity
    prime = 1000003 if C!=4 else 1000033
    FF = GF(prime)
    # zeta0 stands for the C's root of unity
    zeta0 = FF.multiplicative_generator()^((prime-1)/C)

    # In what follows, a point on E is represented by a 
    # pair (X,Y) satisfying X^A-Y^B=1. The zero point is 
    # represented by None
    # We define some functions that are used in some of
    # the experiments:

    # Return a random point on E
    def rand_point():
        PR.<t> = FF[]
        while True:
            Y = FF.random_element()
            pol = t^A-Y^B-1
            rts = pol.roots(FF, multiplicities=False)
            if rts:
                return (rts[0],Y)

    # Return a random point on E^N
    def rand_tuple():
        return [rand_point() for i in range(N)]

    # add 2 points pt1, pt2 on E
    def elladd(pt1, pt2):
        return point2pair(pair2point(pt1)+pair2point(pt2))

    # multiply a point pt by an integer k
    def ellmul(pt, k):
        return point2pair(pair2point(pt)*ZZ(k))

    # apply the automorphism to pt
    def ellzeta(pt):
        if pt is None:
            return None
        return (pt[0]/zeta0^A2, pt[1]/zeta0^B2)

    # compute the sum $\sum_i m_i p_i$ for $p_i\in E$,
    # $m_i\in Z[\zeta]$. The arguments are lists 
    # pts=(p_1,...,p_k), vec=(m_1,...,m_k)
    def ellmulvec(pts, vec):
        res = None
        for cf, pt in zip(vec, pts):
            cf = CF(cf)
            res = elladd(res, ellmul(pt, cf[0]))
            res = elladd(res, ellmul(ellzeta(pt), cf[1]))
        return res

    # find a random point on a divisor on E^N represented by
    # a vector 'divisor' of elements of CF
    def rand_tuple_div(divisor):
        divisor = [CF(x) for x in divisor]
        # find an invertible entry in the vector
        pivot = 0
        while divisor[pivot].norm()!=1:
            pivot += 1
        piv_val = divisor[pivot]
        # normalize the vector so that divisor[pivot]=-1
        for i in range(N):
            divisor[i] /= -piv_val
        # consruct a random point
        randpts = rand_tuple()
        # compute the value at position pivot
        pt = None
        for i in range(N):
            if i==pivot:
                continue
            pt = elladd(pt, ellmul(randpts[i], divisor[i][0]))
            pt = elladd(pt, ellmul(ellzeta(randpts[i]), divisor[i][1]))
        assert pt[0]^A==pt[1]^B+1
        randpts[pivot] = pt
        return randpts

    # pick a random point
    randpts = rand_tuple()
    # extract the x coordinates
    xx = [x for x, _ in randpts]
    # extract the y coordinates
    yy = [y for _, y in randpts]
    # the working field is the finite field
    RF = FF
else:
    # Here the working field contains the coordinates of the point as variables
    xnames, ynames = [[f % (i,) for i in range(N)] for f in ['x%d', 'y%d']]
    names = xnames+ynames+['z']
    order = TermOrder('wdegrevlex', [B]*N+[A]*N)+TermOrder('degrevlex',1)
    RF0 = PolynomialRing(QQ, names=names, order=order)
    xx = list(RF0.gens()[:N])
    yy = list(RF0.gens()[N:2*N])
    zeta0 = RF0.gen(2*N)
    RF = RF0
    FF = RF
    Reqns = [zeta.minpoly()(zeta0)] + [xx[i]^A-yy[i]^B-1 for i in range(N)]
    def myreduce(f):
        return f.reduce(Reqns)
    z2red = (zeta0^2).reduce(Reqns)

    # SAGE does not want to compute gcd of multivariable rational functions properly
    # here is our solution:
    def gcdZ(lst):
        lcmden = lcm([x.denominator() for x in lst])
        return gcd([x*lcmden for x in lst])/lcmden
    # gcd of a pair
    def mygcd(f1, f2):
        gcd0 = f1.gcd(f2)
        gcdcoeff = gcdZ((gcdZ(f1.coefficients()), gcdZ(f2.coefficients())))
        gcdgcd = gcdZ(gcd0.coefficients())
        return gcd0*(gcdcoeff/gcdgcd)
    # gcd of a list
    def mygcdlist(lst):
        assert len(lst)>0
        srt = [(f.number_of_terms(), f) for f in lst]
        srt.sort()
        lst = [f for _, f in srt]
        res = lst[0]
        for f in lst[1:]:
            if res==1:
                return res
            res = mygcd(res, f)
        return res
    # f1, f2 -> (f1/g, g, f2/g) where g=gcd(f1,f2)
    def splitfacs(f1, f2):
        g = mygcd(f1, f2)
        q1, r1 = f1.quo_rem(g)
        q2, r2 = f2.quo_rem(g)
        assert r1.is_zero() and r2.is_zero()
        return (q1, g, q2)

# the matrices from the paper:
H1296 = matrix(CF, 4, 4, [6, 0, 0, 0, \
                0, 6, 0, 0, \
                0, 0, 6, 0, \
                0, 0, 0, 6])
H864 = matrix(CF, 4, 4, [6, 2*zeta-4, 0, 0, \
                2*zeta-4, 6, 0, 0, \
                0, 0, 6, 0, \
                0, 0, 0, 6])
H576 = matrix(CF, 4, 4, [6, 2*zeta-4, 0, 0, \
                2*zeta-4, 6, 0, 0, \
                0, 0, 6, 2*zeta-4, \
                0, 0, 2*zeta-4, 6])
H432 = matrix(CF, 4, 4, [6, 2*zeta-4, 0, 0, \
                2*zeta-4, 6, 2*zeta-4, 0, \
                0, 2*zeta-4, 6, 0, \
                0, 0, 0, 6])
H144 = matrix(CF, 4, 4, [6, 2*zeta - 4, 0, 0, \
                -2*zeta - 2, 6, 2*zeta - 4, 0, \
                0, -2*zeta - 2, 6, 2*zeta - 4, \
                0, 0, -2*zeta - 2, 6])
H13824 = matrix(CF, 5, 5, [12, 4*zeta + 4, 4*zeta + 4, 4*zeta + 4, 4*zeta + 4, \
    -4*zeta + 8, 12, 4*zeta + 4, 4*zeta + 4, 4*zeta + 4, \
    -4*zeta + 8, -4*zeta + 8, 12, 4*zeta + 4, 6, \
    -4*zeta + 8, -4*zeta + 8, -4*zeta + 8, 12, 6, \
    -4*zeta + 8, -4*zeta + 8, 6, 6, 12])

# in each of the 3 cases we have a function Mtoh, which converts
# the discrete invariant given by a matrix M with given row and column sums
# to the corresponding off-diagonal entry of the matrix H 
# myH is the H-matrix that will be used for the computation
if (A,B,C)==(2,4,4):
    Mtoh = lambda M: K*(1-zeta)-M[0,0]-M[1,2] + zeta*(M[0,1]+M[1,3])
    if N==2:
        myH = matrix(CF, N, N, [2, zeta-1, -zeta-1, 2])
    elif N==3:
        myH = matrix(CF, N, N, [4, 2, 2*zeta, 2, 4, 2, -2*zeta, 2, 4])
if (A,B,C)==(2,3,6):
    Mtoh = lambda M: (4-2*zeta)*K-(2+2*zeta)*M[0,0]+(4*zeta-2)*M[1,2]
    EC = EllipticCurve(RF, [0,1])
    def pair2point(pt):
        return EC((pt[1],pt[0])) if pt is not None else EC(0)
    def point2pair(pt):
        return (pt[1],pt[0]) if pt!=EC(0) else None
    if N==4:
        myH = H144
    elif N==5:
        myH = H13824
if (A,B,C)==(3,3,3):
    Mtoh = lambda M: 2 + 2*zeta + (zeta^2*(M[1,1]+M[0,2]+M[2,0])-(M[0,0]+M[1,2]+M[2,1]))/(1-zeta)

# returns all n1\times n2 matrices with non-negative integer entries
# with row sums s1 and column sums s2
def matrices(n1, n2, s1, s2):
    res = []
    mx = min(s1, s2)
    comps = list(Compositions(s1+n2, length=n2))
    vectors = [[x-1 for x in comp] for comp in comps]
    for aa in cartesian_product_iterator([vectors]*(n1-1)):
        if any(sum(aa[i][j] for i in range(n1-1))>s2 for j in range(n2)):
            continue
        m = matrix(QQ, n1-1, n2, aa)
        v = [s2-sum(aa[i][j] for i in range(n1-1)) for j in range(n2)]
        m = m.stack(vector(v))
        yield m

# Ms is the list of all possible M-matrices aka discrete invariants
Ms = list(matrices(A, B, A1, B1))
#print len(Ms)

possible_vals = set()
scal = {}
scal_to_M = {}
allM = []

pairs = [(i,j) for i,j in mrange((N, N)) if i<j]
pairs2 = [(i,j) for i,j in mrange((A, B))]

N2 = len(pairs)

counts = set()

# Useful functions to convert between M-matrices and H-matrices 
# (discrete invariants)
def htoMs(h):
    return [M for M in Ms if Mtoh(M)==h]

def HtoMMs(H):
    return list(cartesian_product_iterator([htoMs(H[i,j]) for i,j in pairs]))

def MMtoH(MM):
    H = matrix(CF, N, N)
    for i in range(N):
        H[i,i] = K*C^2/A/B
    for (i,j), M in zip(pairs, MM):
        H[i,j] = Mtoh(M)
        H[j,i] = H[i,j].conjugate()
    return H

def findMM(H):
    for MM in cartesian_product_iterator([Ms]*N2):
        if MMtoH(MM)==H:
            yield MM

# This function is used to create the base ring for computations
# f is the base field
# each entry in names corresponds to a block of variables of length
# given by the corresponding entry of orders
# for instance if names=['x','y','z'] and orders=[2,3,4], we will have variables
# x2,x1,y3,y2,y1,z4,z3,z2,z1 (each block is sorted backwards)
# the order is the weighted degrevlex order where xi has weight i
# names0 allows to create extra single variables 
# also if an entry in orders is -1, it means to create a single variable
# of very large weight.
# Returns a triple:
# ring, list of lists of variables, list of the extra variables
def ring_blocks(f, names, orders, names0, order0):
    resnames = []
    weights = []
    for n, ord in zip(names, orders):
	resnames.extend([n+str(ord-i) for i in range(ord)] if ord>=0 else [n])
	#weights.extend([1 for i in range(ord)] if ord>=0 else [100])
	weights.extend([(ord-i) for i in range(ord)] if ord>=0 else [100])
    #ring = PolynomialRing(f, resnames+names0, order=TermOrder('wdegrevlex', [1]*len(weights)))
    ring = PolynomialRing(f, resnames+names0, order=TermOrder('wdegrevlex', weights))
    #ring = PolynomialRing(f, resnames+names0, order=TermOrder('wdegrevlex', weights)+order0)
    vars = []
    cnt = 0
    for ord in orders:
        if ord<0:
            assert ord==-1
            ord = 1
	vars.append([ring.gen(cnt+i) for i in range(ord)])
	cnt += ord
    return ring, vars, list(ring.gens()[len(weights):])

# Given a sequence coeffs and a variable x produces a monic polynomial
# with given coefficients
def make_poly(coeffs, x):
    return x^len(coeffs) + sum(x^i*cf for i, cf in enumerate(coeffs))

# Produces a solution of P1^A-Q1^B=P2^A-Q2^B from a matrix of gcds
# in gg passing through the point pt1, pt2. This is done by
# solving the "main equation" eq. (2)
# returns the polynomials P1, P2, Q1, Q2 and a list of
# equations eqns that must be satisfied to make this return valid
def reconstruct_pq(gg, pt1, pt2, zeta0):
    Ptemp = [prod(gg[i,j] for j in range(B)) * (pt1[0]-zeta0^(i*C/A)*pt2[0]) for i in range(A)]
    Qtemp = [prod(gg[j,i] for j in range(A)) * (pt1[1]-zeta0^(i*C/B)*pt2[1]) for i in range(B)]
    coeff1 = 1/FF((zeta^(C/A)-1).norm()) * (zeta0^(C-C/A)-1)
    P1 = (Ptemp[0]*zeta0^(C/A)-Ptemp[1]) * coeff1
    P2 = P1-Ptemp[0]
    coeff2 = 1/FF((zeta^(C/B)-1).norm()) * (zeta0^(C-C/B)-1)
    Q1 = (Qtemp[0]*zeta0^(C/B)-Qtemp[1]) * coeff2
    Q2 = Q1-Qtemp[0]
    eqns = [P1-zeta0^(i*C/A)*P2-Ptemp[i] for i in range(2,A)]
    eqns += [Q1-zeta0^(i*C/B)*Q2-Qtemp[i] for i in range(2,B)]
    eqns = sum([eq.coefficients() for eq in eqns], [])
    return (P1, P2, Q1, Q2, eqns)


@cached_function
def cached_resultant(P, Q):
    return P.resultant(Q)

# Our equations are homogeneous and invariant under translations
# To rigidify we set one variable of degree 1 to zero and another one to 1.
# Given the list of lists of variables, returns the resulting equations
# There is a freedom in the choice of degree 1 variables, so we choose those
# which are in a shortest block
def cell_proj_eqns(cfs):
    srt = []
    for j, lst in enumerate(cfs):
        d = len(lst)
        for i, var in enumerate(lst):
            srt.append((d-i, d, j, var))
    srt.sort()
    vars = [var for _, _, _, var in srt]
    #print vars
    eqns = [vars[0], vars[1]-1]
    return eqns

# This is the main computation
# Input: MM contains the M-matrix for each pair of indices
# Remember that M-matrix contains degrees of the gcd's
# Returns ring and an list of 3 lists of equations. Solutions 
# to all the equations are the curves passing through the given point.
# The coordinates of the point are given in the global variables xx, yy
def countsols_pt_eqns(MM):
    res = []
    # Compile a list of opposed pairs, which are pairs
    # of polynomials have to be relatively prime
    # for each such pair g1, g2 we will have a variable inv
    # and an equation inv*resultant(g1, g2)-1=0
    opposed_pairs = []
    for M in MM:
        opp = []
        for k1, (i1, j1) in enumerate(pairs2):
            if M[i1,j1]==0:
                continue
            for k2, (i2, j2) in enumerate(pairs2):
                if M[i2,j2]==0:
                    continue
                if i1<i2 and j1!=j2:
                    opp.append((k1, k2))
        opposed_pairs.append(opp)
    # number of inv-variables
    nopp = sum(len(opp) for opp in opposed_pairs)# + N
    #print nopp, opposed_pairs
    names0 = []
    #names0 = ['x%d'%(i,) for i in range(N)]
    #names0 = ['x%d'%(i,) for i in range(N)] \
    #        + ['y%d'%(i,) for i in range(N)]
    #order0 = TermOrder('degrevlex', N)
    order0 = []
    # for each polynomial g we have a block of variables
    # storing its coefficients
    names = ['c%d'%(i,) for i in range(A*B*N2)]
    degrees = sum([M.list() for M in MM], [])
    # create the ring and the variables
    ring, cfs, var0 = ring_blocks(RF, \
            names+['inv%d'%(i,) for i in range(nopp)], \
            degrees+[-1]*nopp, names0, order0)
    #xx = var0[:N]
    #yy = var0[N:2*N]
    pts = list(zip(xx, yy))
    #zeta0 = var0[-1]
    invs = cfs[-nopp:]
    cfs = cfs[:-nopp]
    ring_T.<T> = ring[]
    Ps = [None]*N
    Qs = [None]*N

    pairsdic = {(i,j):k for k, (i,j) in enumerate(pairs)}
    invcnt = 0
    extraeqns = []
    eqns = []
    #eqns.append(zeta.minpoly(zeta0))
    #for i in range(N):
    #    eqns.append(xx[i]^A-yy[i]^B-1)
    # First come the rigidifying equations
    eqns.extend(cell_proj_eqns(cfs))
    res.append(eqns)
    eqns = []
    # next we compute the Ps and the Qs out of the gs,
    # collecting all the equations that come up into a long array eqns
    # the equations with 'inv' variables, see above, are collected into
    # another array extraeqns
    for j in range(1, N):
        for i in range(j):
            k = pairsdic[i,j]
            gpolys = [make_poly(cf, T) for cf in cfs[k*A*B:(k+1)*A*B]]
            ggM = matrix(ring_T, A, B, gpolys)
            # Compute the Ps and the Qs, and store the equations that came up
            P1v, P2v, Q1v, Q2v, eqns1 = reconstruct_pq(ggM, pts[i], pts[j], zeta0)
            eqns.extend(eqns1)
            # compute the equations inv*resultant(g1,g2)-1=0
            for k1, k2 in opposed_pairs[k]:
                extraeqns.append(invs[invcnt][0]*gpolys[k1].resultant(gpolys[k2])-1)
                invcnt += 1
            # Initially Ps[i] and Qs[i] are not known, so they are set to None
            # In the process, if Ps[i] say is not known, we simply store the computed
            # value. Otherwise, if we already have a value for Ps[i] we have to
            # add the equations saying that the new value must be equal to the old one
            if Ps[i] is None:
                Ps[i] = P1v
            else:
                eqns.extend((Ps[i]-P1v).coefficients())
            if Ps[j] is None:
                Ps[j] = P2v
            else:
                eqns.extend((Ps[j]-P2v).coefficients())
            if Qs[i] is None:
                Qs[i] = Q1v
            else:
                eqns.extend((Qs[i]-Q1v).coefficients())
            if Qs[j] is None:
                Qs[j] = Q2v
            else:
                eqns.extend((Qs[j]-Q2v).coefficients())

    res.append(eqns)
    res.append(extraeqns)
    eqns = []

    #for i in range(N):
    #    eqns.append(invs[invcnt][0]*cached_resultant(Ps[i], Qs[i])-1)
    #    invcnt += 1
    #res.append(eqns)
    
    return ring, res

# These functions are used to find linear relations between points
# on E
def HHtoLattice(H):
    N = H.ncols()
    res = matrix(ZZ, 2*N, 2*N)
    for i in range(2*N):
        for j in range(2*N):
            mul = 1
            if j>=N:
                mul*=zeta
            if i>=N:
                mul*=zeta.conjugate()
            res[i,j] = CF(H[i%N,j%N]*mul).trace()
    return res

def HHvectors_by_length(H, ln):
    N = H.ncols()
    qf = QuadraticForm(HHtoLattice(H))
    res0 = qf.vectors_by_length(ln)[ln]
    res = []
    for v0 in res0:
        v = vector(CF, [v0[i]+v0[i+N]*zeta for i in range(N)])
        res.append(v)
    return res

# Returns all small relations between points given in a list pt 
# small means that all coefficients have norm at most ln
def find_relations(pt, ln):
    pt = [(RF(x),RF(y)) for x,y in pt]
    rels = []
    for v in HHvectors_by_length(myH, ln):
        if ellmulvec(pt, v) is None:
            rels.append(v)
    return rels

# Now the main computation
# each time we get a list of 3 lists of equations 
# we first compute the groebner basis for the equations in 
# the first and the second lists. Then we add the equations
# from the third list and compute the groebner basis of the result
from time import time
import sys
if FINITE_FIELD:
    good_count = 0
    # This will contain for each pair (dim, vdim) how many tests
    # produced this result.
    statistics = {}
    for count in range(NTRIES):
        randpts = rand_tuple() if DIVISOR==None else rand_tuple_div(DIVISOR)
        xx = [x for x, _ in randpts]
        yy = [y for _, y in randpts]
        print randpts
        

        t0 = time()
        ring,eqns = countsols_pt_eqns(HtoMMs(myH)[0])
        #print 'Time:', time()-t0

        #t0 = time()
        J1=ring.ideal(eqns[0]+eqns[1])
        #print J1.dimension()
        #print [gb.lt() for gb in J1.groebner_basis()]
        #print 'Time:', time()-t0

        #t0 = time()
        # Compute the groebner basis for the equations 
        # in the first and the second lists and add the
        # equations from the third list
        J2=ring.ideal(J1.groebner_basis()+eqns[2])
        dim = J2.dimension()
        vdim = None
        if dim==0:
            vdim = J2.vector_space_dimension()
        print dim, vdim
        if (dim,vdim) not in statistics:
            statistics[dim,vdim]=0
        statistics[dim,vdim]+=1
        # we have dim=Krull dimension and vdim=vector space dimension if dim=0
        if dim==0 and vdim==1:
            # this is good
            good_count += 1
        #print [gb.lt() for gb in J2.groebner_basis()]
        print 'Time:', time()-t0
    print 'Result: ', good_count, 'out of', NTRIES
    print 'Statistics (dim,vdim): ', statistics
    sys.stdout.flush()
else:
    pass