#! /usr/local/bin/sage
### -*-python-*-

from itertools import count, izip
from random import randrange

## Constants.
EC_YBIT = EC_XONLY = 1
EC_CMPR = EC_LSB = 2
EC_EXPLY = 4
EC_SORT = 8

## Conversion primitives.
def poly_to_int(x, p):
  return ZZ(sum([ZZ(a)*p**i for a, i in izip(x.list(), count())]))
def int_to_poly(n, p, t):
  return sum([a*t**i for a, i in izip(n.digits(p), count())])

def fe2ip(x):
  return poly_to_int(x.polynomial(), x.parent().characteristic())
def i2fep(n, k):
  return int_to_poly(n, k.characteristic(), k.gen())

def fe2osp(x):
  k = x.parent()
  n = fe2ip(x)
  nb = ceil((k.cardinality() - 1).nbits()/8)
  v = n.digits(256, padto = nb); v.reverse(); return v

def os2sp(v):
  return ''.join('%02x' % i for i in v)

def fe2sp(x):
  if x == 0 or x == 1:
    return '%d' % x
  if x.parent().degree() > 1:
    return '0x%x' % fe2ip(x)
  else:
    n = ZZ(x + 99)
    if n < 199: return '%d' % (n - 99)
    else: return '0x%x' % x

def os_to_hex(v):
  return ''.join('%02x' % b for b in v)

## Other utilities.
def find_2torsion_point(E, n):
  while is_even(n): n //= 2
  while True:
    Q = E.random_point()
    R = n*Q
    if R: break
  while True:
    S = 2*R
    if not S: break
    R = S
  return R

def pick_at_random(choices):
  return choices[randrange(len(choices))]

## Output an elliptic curve description.
def ecdescr(E):
  k = E.base_ring()
  m = k.degree()
  if m == 1:
    fdesc = '%s: %d' % (pick_at_random(['prime', 'niceprime']),
                        k.characteristic())
    cdesc = '%s: %s, %s' % (pick_at_random(['prime', 'primeproj']),
                            fe2sp(E.a4()), fe2sp(E.a6()))
  elif k.characteristic() == 2:
    fdesc = '%s: 0x%x' % ('binpoly', poly_to_int(k.polynomial(), 2))
    cdesc = '%s: %s, %s' % (pick_at_random(['bin', 'binproj']),
                            fe2sp(E.a2()), fe2sp(E.a6()))
  else:
    raise ValueError, 'only prime or binary fields supported'
  return '"%s; %s"' % (fdesc, cdesc)

## Output a point description.
def ecpt(Q):
  if Q is None: return 'FAIL'
  elif not Q: return 'inf'
  else: return '"%s, %s"' % (fe2sp(Q[0]), fe2sp(Q[1]))

## Compress a point.
def ptcompr_lsb(Q):
  x, y = Q[0], Q[1]
  if Q.curve().base_ring().characteristic() != 2:
    return is_odd(fe2ip(y)) and EC_YBIT or 0
  elif x == 0:
    return 0
  else:
    return (y/x).polynomial()[0] == 1 and EC_YBIT or 0

def ptcompr_sort(Q):
  y = fe2ip(Q[1]); yy = fe2ip((-Q)[1])
  return y > yy and EC_YBIT or 0

## Some elliptic curves.  Chosen to have 2-torsion, so as to test edge cases
## where Q = -Q; in particular, 2-torsion points can't have their
## y-coordinates compressed to 1.
p = previous_prime(2^192); k_p = GF(p)
E_p = EllipticCurve(k_p, [-3, 6])
n_p = 6277101735386680763835789423293484020607766684840576538738
T_p = find_2torsion_point(E_p, n_p)

k_2.<t> = GF(2^191)
E_2 = EllipticCurve(k_2, [1, 1, 0, 0, 1])
n_2 = 3138550867693340381917894711593325610042432240066118385966
T_2 = find_2torsion_point(E_2, n_2)

def test_ec2osp(E, Q):
  ec = ecdescr(E)
  pt = ecpt(Q)
  xs, ys = fe2osp(Q[0]), fe2osp(Q[1])
  ybit_lsb = ptcompr_lsb(Q)
  ybit_sort = ptcompr_sort(Q)

  def out(f, s):
    print '  %s\n    %d %s\n    %s;' % (ec, f, pt, os_to_hex(s))

  out(EC_XONLY, [EC_XONLY] + xs)
  out(EC_CMPR, [EC_CMPR | ybit_lsb] + xs)
  out(EC_CMPR | EC_SORT, [EC_CMPR | EC_SORT | ybit_sort] + xs)
  out(EC_EXPLY, [EC_EXPLY] + xs + ys)
  out(EC_CMPR | EC_EXPLY, [EC_CMPR | EC_EXPLY | ybit_lsb] + xs + ys)
  out(EC_CMPR | EC_SORT | EC_EXPLY,
      [EC_CMPR | EC_SORT | EC_EXPLY | ybit_sort] + xs + ys)

print 'ec2osp {'
for i in xrange(20): test_ec2osp(E_p, E_p.random_point())
test_ec2osp(E_p, T_p)
for i in xrange(20): test_ec2osp(E_2, E_2.random_point())
test_ec2osp(E_2, T_2)
print '}'

def test_os2ecp(E, Q):
  ec = ecdescr(E)
  k = E.base_ring()
  m = k.degree()
  x, y = Q[0], Q[1]
  xs, ys = fe2osp(Q[0]), fe2osp(Q[1])
  ybit_lsb = ptcompr_lsb(Q)
  ybit_sort = ptcompr_sort(Q)

  def out(f, s, Q):
    sufflen = randrange(16)
    suff = [randrange(256) for i in xrange(sufflen)]
    print '  %s\n    %d\n    %s\n    %s\n    %d;' % (
      ec, f, os_to_hex(s + suff), ecpt(Q), sufflen)

  ## This is the algorithm from `gfreduce_quadsolve'.
  B = x^3 + E.a2()*x^2 + E.a4()*x + E.a6()
  A = E.a1()*x + E.a3()
  if A == 0:
    yy = sqrt(B)
  else:
    xx = B/A^2
    while True:
      rho = k.random_element()
      if rho.trace() == 1: break
    z = sum(rho^(2^i)*sum(xx^(2^j) for j in xrange(i)) for i in xrange(m))
    if z.polynomial()[0]: z -= 1
    yy = A*z

  out(EC_XONLY, [EC_XONLY] + xs, E(x, yy))
  out(EC_EXPLY, [EC_EXPLY] + xs + ys, Q)

  out(EC_LSB, [EC_CMPR | ybit_lsb] + xs, Q)
  out(EC_LSB | EC_EXPLY, [EC_EXPLY | EC_CMPR | ybit_lsb] + xs + ys, Q)
  out(EC_LSB | EC_EXPLY,
      [EC_EXPLY | EC_CMPR | (not ybit_lsb)] + xs + ys, None)

  out(EC_SORT, [EC_CMPR | EC_SORT | ybit_sort] + xs, Q)
  out(EC_SORT | EC_EXPLY,
      [EC_EXPLY | EC_CMPR | EC_SORT | ybit_sort] + xs + ys, Q)
  out(EC_SORT | EC_EXPLY,
      [EC_EXPLY | EC_CMPR | EC_SORT | (not ybit_sort)] + xs + ys, None)

def test_os2ecp_2torsion(E, Q):
  ec = ecdescr(E)
  k = E.base_ring()
  m = k.degree()
  x, y = Q[0], Q[1]
  xs, ys = fe2osp(Q[0]), fe2osp(Q[1])

  def out(f, s, Q):
    sufflen = randrange(16)
    suff = [randrange(256) for i in xrange(sufflen)]
    print '  %s\n    %d\n    %s\n    %s\n    %d;' % (
      ec, f, os_to_hex(s + suff), ecpt(Q), sufflen)

  out(EC_LSB, [EC_CMPR] + xs, Q)
  out(EC_LSB, [EC_CMPR | EC_YBIT] + xs, None)
  out(EC_SORT, [EC_SORT | EC_CMPR] + xs, Q)
  out(EC_SORT, [EC_SORT | EC_CMPR | EC_YBIT] + xs, None)

def test_os2ecp_notpoints(E):
  ec = ecdescr(E)
  k = E.base_ring()

  def out(f, s):
    print '  %s\n    %d\n    %s\n    FAIL\n    0;' % (ec, f, os_to_hex(s))

  while True:
    x = k.random_element()
    if not E.is_x_coord(x): break
  xs = fe2osp(x)
  out(EC_XONLY, [EC_XONLY] + xs)
  out(EC_LSB, [EC_CMPR] + xs)
  out(EC_SORT, [EC_CMPR | EC_SORT] + xs)

print 'os2ecp {'
for i in xrange(20): test_os2ecp(E_p, E_p.random_point())
test_os2ecp_2torsion(E_p, T_p)
test_os2ecp_notpoints(E_p)
for i in xrange(20): test_os2ecp(E_2, E_2.random_point())
test_os2ecp_2torsion(E_2, T_2)
test_os2ecp_notpoints(E_2)
print '}'