#! /usr/bin/python

from sys import argv
import catacomb as C

TESTS = {}
NTEST = 20

def test(arg):
  def reg(fn, name):
    TESTS[name] = fn
    return fn
  if isinstance(arg, str): return lambda fn: reg(fn, arg)
  else: return reg(arg, arg.__name__.replace('_', '-'))

FIELDS = {}

class FieldElt (object):
  def __init__(me, k, n):
    me.k = k
    me.v = (C.MP(n)%k.p).storel(k.len)
  def __str__(me): return hex(me.v)
  @property
  def n(me): return C.MP.loadl(me.v)
  def __pos__(me): return FieldElt(me.k, me.n)
  def __neg__(me): return FieldElt(me.k, -me.n)
  def __nonzero__(me): return me.n != 0
  def __add__(me, you): return FieldElt(me.k, me.n + me.k(you).n)
  def __radd__(me, you): return FieldElt(me.k, me.k(you).n + me.n)
  def __sub__(me, you): return FieldElt(me.k, me.n - me.k(you).n)
  def __rsub__(me, you): return FieldElt(me.k, me.k(you).n - me.n)
  def __mul__(me, you): return FieldElt(me.k, me.n*me.k(you).n)
  def __rmul__(me, you): return FieldElt(me.k, me.k(you).n*me.n)
  def __div__(me, you): return FieldElt(me.k, me.n*me.k(you).inv().n)
  def __rdiv__(me, you): return FieldElt(me.k, me.k(you).n*me.inv().n)
  def inv(me): return FieldElt(me.k, me.k.p.modinv(me.n))
  def sqrt(me): return FieldElt(me.k, me.k.p.modsqrt(me.n))
  @classmethod
  def rand(cls, k):
    me = cls(k, 0)
    me.v = C.rand.block(k. len)
    return me

class Field (object):
  def __init__(me, p, len = None):
    me.p = C.MP(p)
    me.len = len is None and me.p.noctets or len
  @classmethod
  def register(cls, name, *args, **kw):
    FIELDS[name] = cls(*args, **kw)
  def rand(me): return FieldElt.rand(me)
  def __call__(me, n):
    if isinstance(n, FieldElt):
      assert n.k is me
      return n
    else:
      return FieldElt(me, n)

Field.register('f25519', C.MP(0).setbit(255) - 19)
Field.register('fgoldi', C.MP(0).setbit(448) - C.MP(0).setbit(224) - 1)

def binop(k, op):
  x = k.rand(); y = k.rand()
  print '  %s\n    %s\n    %s;' % (x, y, op(x, y))

def unop(k, op):
  x = k.rand()
  print '  %s\n    %s;' % (x, op(x))

@test
def add(k): binop(k, lambda x, y: x + y)

@test
def sub(k): binop(k, lambda x, y: x - y)

@test
def neg(k): unop(k, lambda x: -x)

@test
def mul(k): binop(k, lambda x, y: x*y)

@test
def sqr(k): unop(k, lambda x: x*x)

@test
def inv(k): unop(k, lambda x: x and x.inv() or k(0))

@test
def quosqrt(k):
  x = k.rand(); y = k.rand()
  yy = +y
  u = yy and x/y or k(0)
  try: zz = u.sqrt()
  except ValueError: print '  %s\n    %s\n    "" "";' % (x, y)
  else: print '  %s\n    %s\n    %s\n    %s;' % (x, y, zz, -zz)

@test
def mulconst(k):
  x = k.rand()
  a = C.rand.range(1 << 20) - (1 << 19)
  print '  %s %d\n    %s;' % (x, a, a*x)

def mask(): return C.rand.range(2)*0xffffffff

@test
def pick2(k):
  x = k.rand(); y = k.rand(); m = mask()
  print '  %s\n    %s\n    0x%08x\n    %s;' % (x, y, m, +(m and x or y))

@test
def condneg(k):
  x = k.rand(); m = mask()
  print '  %s\n    0x%08x\n    %s;' % (x, m, x*(m and -1 or +1))

@test
def pickn(k):
  n = C.rand.range(31) + 1
  v = [k.rand() for i in xrange(n)]
  i = C.rand.range(n)
  print '  "%s"\n    %d\n    %s;' % ('\n   '.join(map(str, v)), i, +v[i])

@test
def condswap(k):
  x = k.rand(); y = k.rand(); m = mask()
  xx, yy = m and (+y, +x) or (+x, +y)
  print '  %s\n    %s\n    0x%08x\n    %s\n    %s;' % (x, y, m, xx, yy)

@test
def sub_mulc_add_sub_mul(k):
  u = k.rand(); v = k.rand(); w = k.rand(); x = k.rand(); y = k.rand();
  a = C.rand.range(1 << 20) - (1 << 19)
  print '  %s\n    %s %d\n    %s\n    %s\n    %s\n    %s;' % \
      (u, v, a, w, x, y, ((u - v)*a + w)*(x - y))

k = FIELDS[argv[1]]
for t in argv[2:]:
  print '%s {' % t
  for i in xrange(NTEST): TESTS[t](k)
  print '}'