import sys
import random

print('version 2020.03.17')

# biquadratics for montgomery model (with B=1)

for loop in range(1000):
  q = random.choice([7,9,11,13,17,19,23,25,27,29])
  k.<fieldgen> = GF(q)
  kX.<X> = k[]

  while True:
    A = k.random_element()
    a1 = 0
    a2 = A
    a3 = 0
    a4 = 1
    a6 = 0
    try:
      E = EllipticCurve([a1,a2,a3,a4,a6])
      if E.cardinality().is_power_of(2): continue
      break
    except:
      pass

  while True:
    P = E.random_point()
    Q = E.random_point()
    if P == 0: continue
    if Q == 0: continue
    if P == Q: continue
    if P == -Q: continue
    break

  def F0(X1,X2):
    return (X1-X2)^2

  def F1(X1,X2):
    return -2*((X1*X2+1)*(X1+X2)+2*A*X1*X2)

  def F2(X1,X2):
    return (X1*X2-1)^2

  print('try','biquadratic',q,A,P,Q)
  sys.stdout.flush()

  # check that biquadratics work:

  xP = P[0]
  xQ = Q[0]
  xplus = (P+Q)[0]
  xminus = (P-Q)[0]
  assert (X-xplus)*(X-xminus) == X^2+X*F1(xP,xQ)/F0(xP,xQ)+F2(xP,xQ)/F0(xP,xQ)

  # triquadratic:

  kZ0Z1Z2.<Z0,Z1,Z2> = PolynomialRing(k,3)
  assert (Z0*Z1-1)^2+(Z0*Z2-1)^2+(Z1*Z2-1)^2-2*Z0*Z1*Z2*(Z0+Z1+Z2+2*A)-2 == Z0^2*F0(Z1,Z2)+Z0*F1(Z1,Z2)+F2(Z1,Z2)

# Algorithm 2

for loop in range(100):
  q = random.choice([7,9,11,13,17,19,23,25,27,29])
  k.<fieldgen> = GF(q)
  kX.<X> = k[]

  while True:
    A = k.random_element()
    a1 = 0
    a2 = A
    a3 = 0
    a4 = 1
    a6 = 0
    try:
      E = EllipticCurve([a1,a2,a3,a4,a6])
      if E.cardinality().is_power_of(2): continue
      break
    except:
      pass

  while True:
    P = E.random_point()
    n = P.order()
    if n >= 3: break

  I = set()
  J = set()
  K = set()
  S = set()
  for ijloop in range(loop):
    if randrange(2):
      inew = randrange(-loop*loop,loop*loop+1)
      if not inew%n: continue
      if inew in I: continue
      if any(inew+j in S for j in J): continue
      if any(inew-j in S for j in J): continue
      if not all((inew+j)%n for j in J): continue
      if not all((inew-j)%n for j in J): continue
      I.add(inew)
      for j in J:
        S.add(inew+j)
        S.add(inew-j)
    else:
      jnew = randrange(1,loop*loop+1)
      if not jnew%n: continue
      if jnew in J: continue
      if any(i+jnew in S for i in I): continue
      if any(i-jnew in S for i in I): continue
      if any(i+2*jnew in I for i in I): continue
      if not all((i+jnew)%n for i in I): continue
      if not all((i-jnew)%n for i in I): continue
      J.add(jnew)
      for i in I:
        S.add(i+jnew)
        S.add(i-jnew)

  for ijloop in range(loop):
    snew = randrange(-loop*loop,loop*loop+1)
    if not snew%n: continue
    if snew in S: continue
    K.add(snew)
    S.add(snew)

  assert all(s%n for s in S)
  assert all(i%n for i in I)
  assert all(j%n for j in J)
  IJ = set(i+j for i in I for j in J).union(set(i-j for i in I for j in J))
  assert IJ.isdisjoint(K)
  assert IJ.union(K) == S
  assert len(IJ) == 2*len(I)*len(J)
  assert len(S) == 2*len(I)*len(J)+len(K)

  alpha = k.random_element()

  print('try','algorithm2',q,A,P,n,alpha,len(S),len(I),len(J),len(K))
  sys.stdout.flush()

  x = {s:(s*P)[0] for s in S.union(I).union(J)}
  hI = kX(prod(X-x[i] for i in I))
  DJ = kX(prod(F0(X,x[j]) for j in J))
  DeltaIJ = hI.resultant(DJ)
  assert DeltaIJ != 0
  EJ = kX(prod(F0(X,x[j])*alpha^2+F1(X,x[j])*alpha+F2(X,x[j]) for j in J))
  R = hI.resultant(EJ)
  hK = prod(alpha-x[s] for s in K)
  assert hK*R/DeltaIJ == prod(alpha-x[s] for s in S)

# evaluating kernel polynomials

for loop in range(200):
  q = random_prime(1000,lbound=100)
  k.<fieldgen> = GF(q)
  kX.<X> = k[]

  while True:
    A = k.random_element()
    a1 = 0
    a2 = A
    a3 = 0
    a4 = 1
    a6 = 0
    try:
      E = EllipticCurve([a1,a2,a3,a4,a6])
      if E.cardinality().is_power_of(2): continue
      break
    except:
      pass

  while True:
    P = E.random_point()
    l = P.order()
    if l%2: break

  S = set(range(1,l-1,2))
  if l == 1:
    assert S == set()
  else:
    assert min(S) == 1
    assert max(S) == l-2
    assert len(S) == (l-1)/2

  b = floor(sqrt(l-1)/2)
  if b == 0:
    b2 = 0
  else:
    b2 = floor((l-1)/(4*b))
  I = set(2*b*(2*i+1) for i in range(b2))
  J = set(2*j+1 for j in range(b))
  if b == 0:
    assert J == set()
  else:
    assert min(J) == 1
    assert max(J) == 2*b-1
  K = range(4*b*b2+1,l-1,2)

  IJ = set(i+j for i in I for j in J).union(set(i-j for i in I for j in J))
  assert IJ.isdisjoint(K)
  assert IJ.union(K) == S
  assert len(IJ) == 2*len(I)*len(J)
  assert len(S) == 2*len(I)*len(J)+len(K)

  print('try','kerpoly',q,A,P,l,alpha,len(S),len(I),len(J),len(K))
  sys.stdout.flush()

  alpha = k.random_element()

  x = {s:(s*P)[0] for s in S.union(I).union(J)}
  hI = kX(prod(X-x[i] for i in I))
  DJ = kX(prod(F0(X,x[j]) for j in J))
  DeltaIJ = hI.resultant(DJ)
  assert DeltaIJ != 0
  EJ = kX(prod(F0(X,x[j])*alpha^2+F1(X,x[j])*alpha+F2(X,x[j]) for j in J))
  R = hI.resultant(EJ)
  hK = prod(alpha-x[s] for s in K)
  assert hK*R/DeltaIJ == prod(alpha-x[s] for s in S)

  h = kX(prod(X-x[s] for s in S))
  assert hK*R/DeltaIJ == h(alpha)

  # derivative

  print('try','kerpolydiff',q,A,P,l,alpha,len(S),len(I),len(J),len(K))
  sys.stdout.flush()

  kjets.<epsilon> = k[[]]
  kjetsY.<Y> = kjets[]

  hI = kjetsY(prod(Y-x[i] for i in I))
  DJ = kjetsY(prod(F0(Y,x[j]) for j in J))
  assert DeltaIJ == hI.resultant(DJ)
  alphajet = alpha+epsilon+O(epsilon^2)
  EJ = kjetsY(prod(F0(Y,x[j])*alphajet^2+F1(Y,x[j])*alphajet+F2(Y,x[j]) for j in J))
  R = hI.resultant(EJ)
  hK = prod(alphajet-x[s] for s in K)
  R /= DeltaIJ
  R *= hK
  assert R[0] == h(alpha)
  assert R[1] == h.diff()(alpha)

# I,J,K example for step-2 arithmetic progression:

for m in range(1,1000,2):
  S = set(range(1,m+2,2))
  assert min(S) == 1
  assert max(S) == m
  b = floor(sqrt(m+1)/2)
  if b == 0:
    b2 = 0
  else:
    b2 = floor((m+1)/(4*b))
  I = set(2*b*(2*i+1) for i in range(b2))
  J = set(2*j+1 for j in range(b))
  K = range(4*b*b2+1,m+1,2)

  print('try','step2',m,len(S),len(I),len(J),len(K))
  sys.stdout.flush()

  IJ = set(i+j for i in I for j in J).union(set(i-j for i in I for j in J))
  assert IJ.isdisjoint(K)
  assert IJ.union(K) == S
  assert len(IJ) == 2*len(I)*len(J)
  assert len(S) == 2*len(I)*len(J)+len(K)
  if b > 0:
    assert len(I) == b2
    assert b2 <= b+2
    assert len(J) == b
    assert len(K) <= 2*b-1

# not used in paper, various options for step-1 arithmetic progression:

for m in range(1000):
  # reasonably efficient option:

  S = set(range(m))
  if m > 0:
    assert min(S) == 0
    assert max(S) == m-1
  b = floor(sqrt(m/2))
  b2 = ceil(b/2)
  if b == 0:
    b3 = 0
  else:
    b3 = max(0,floor((m+b-2*b2)/(2*b)))
  I = set(b2+b*(b3-1)+i for i in range(b))
  J = set(b2+b*j for j in range(b3))
  K = set(range(2*b*b3+2*b2-b,m))
  if b%2 and b*b3 < m: K.add(b*b3)

  print('try','step1',m,len(S),len(I),len(J),len(K))
  sys.stdout.flush()

  IJ = set(i+j for i in I for j in J).union(set(i-j for i in I for j in J))
  assert IJ.isdisjoint(K)
  assert IJ.union(K) == S
  assert len(IJ) == 2*len(I)*len(J)
  assert len(S) == 2*len(I)*len(J)+len(K)
  if b > 0:
    assert len(I) == b
    assert len(J) == b3
    assert b3 <= b+2
    assert len(K) <= 2*b

  # simpler but less efficient option:

  if m > 0:
    b = floor((1+sqrt(8*m-7))/4)
  else:
    b = 0

  I = set(b-1+i*(2*b-1) for i in range(b))
  J = set(range(1,b))
  K1 = I
  K2 = set(range(2*b^2-b,m))
  K = K1.union(K2)
  if m > 0:
    assert 2*b^2-b < m
  assert 2*(b+1)^2-(b+1) >= m

  print('try','step1simpler',m,len(S),len(I),len(J),len(K))
  sys.stdout.flush()

  IJ = set(i+j for i in I for j in J).union(set(i-j for i in I for j in J))
  assert IJ.isdisjoint(K)
  assert IJ.union(K) == S
  assert len(IJ) == 2*len(I)*len(J)
  assert len(S) == 2*len(I)*len(J)+len(K)
  if b > 0:
    assert len(I) == b
    assert len(J) == b-1
    assert len(K1) == b
    assert len(K2) <= 4*b+1

  # simplest but less efficient option:

  b = floor(sqrt(m/2))
  I = set((2*i+1)*b-1 for i in range(b))
  J = set(range(1,b))
  if b == 0:
    K = S
  else:
    K1 = set(i*b-1 for i in range(1,2*b))
    K2 = set(range(2*b^2-1,m))
    K = K1.union(K2)

  print('try','step1simplest',m,len(S),len(I),len(J),len(K))
  sys.stdout.flush()

  IJ = set(i+j for i in I for j in J).union(set(i-j for i in I for j in J))
  assert IJ.isdisjoint(K)
  assert IJ.union(K) == S
  assert len(IJ) == 2*len(I)*len(J)
  assert len(S) == 2*len(I)*len(J)+len(K)
  if b > 0:
    assert len(I) == b
    assert len(J) == b-1
    assert len(K1) == 2*b-1
    assert len(K2) <= 4*b+2