#!/usr/bin/env python3
import random
import gmpapi

MAX_SLONG="9223372036854775807"
MIN_SLONG="-9223372036854775808"
MAX_ULONG="18446744073709551615"
MAX_SINT="2147483647"
MIN_SINT="-2147483648"
MAX_UINT="4294967295"
MAX_SSHORT="32767"
MIN_SSHORT="-32768"
MAX_USHORT="65535"

def plus1(x):  return str(int(x)+1)
def minus1(x): return str(int(x)-1)
def apply(fun, lst):
  return list(map(str, map(fun, lst)))

mzero_one = ["-0", "-1"]
zero_one  = ["0", "1"]
mm_slong   = [MAX_SLONG, MIN_SLONG]
mm_slong1  = [minus1(MAX_SLONG), plus1(MIN_SLONG)]
mm_ulong  = [MAX_ULONG]
mm_ulong1 = [minus1(MAX_ULONG)]
mm_sint   = [MAX_SINT, MIN_SINT]
mm_sint1  = [minus1(MAX_SINT), plus1(MIN_SINT)]
mm_uint   = [MAX_UINT]
mm_uint1  = [minus1(MAX_UINT)]
mm_sshort = [MAX_SSHORT, MIN_SSHORT]
mm_sshort1= [minus1(MAX_SSHORT), plus1(MIN_SSHORT)]
mm_ushort = [MAX_USHORT]
mm_ushort1= [minus1(MAX_USHORT)]
mm_all    = mm_slong + mm_ulong + mm_sint + mm_uint + mm_sshort + mm_ushort
zero_one_all = mzero_one + zero_one


mpz_std_list = zero_one_all + mm_all + apply(plus1, mm_all) + apply(minus1, mm_all)
si_std_list  = zero_one     + mm_slong + mm_sint + mm_sshort + mm_slong1 + mm_sint1 + mm_sshort1
ui_std_list  = zero_one     + mm_ulong + mm_uint + mm_ushort + mm_ulong1 + mm_uint1 + mm_ushort1

def gen_random_mpz(mindigits=1, maxdigits=100, allowneg=True):
  sign = random.choice(["", "-"])
  if not allowneg:
    sign = ""
  return sign + gen_digits(random.randint(mindigits, maxdigits))

def gen_random_si():
  si = gen_random_mpz(mindigits=1,maxdigits=19)
  while int(si) > int(MAX_SLONG) or int(si) < int(MIN_SLONG):
    si = gen_random_mpz(mindigits=1,maxdigits=19)
  return si

def gen_random_ui():
  ui = gen_random_mpz(mindigits=1,maxdigits=20, allowneg=False)
  while int(ui) > int(MAX_ULONG):
    ui = gen_random_mpz(mindigits=1,maxdigits=20, allowneg=False)
  return ui

def gen_digits(length):
  if length == 1:
    i = random.randint(1, 9)
  else:
    digits = [random.randint(1,9)] + [random.randint(0,9) for x in range(length-1)]
    digits = map(str, digits)
    i = "".join(digits)
  return str(i)

def gen_mpzs(mindigits=1, maxdigits=100, count=10):
  return [gen_random_mpz(mindigits=mindigits, maxdigits=maxdigits) for x in range(count)]

default_count = 10
def gen_sis(count=default_count):
  return [gen_random_si() for x in range(count)]

def gen_uis(count=default_count):
  return [gen_random_ui() for x in range(count)]

def gen_small_mpzs(count=default_count):
  return gen_mpzs(mindigits=1,maxdigits=4,count=count)

def is_small_mpz(s):
  return len(s) >= 1 and len(s) <= 4

def gen_medium_mpzs(count=default_count):
  return gen_mpzs(mindigits=5,maxdigits=20,count=count)

def is_medium_mpz(s):
  return len(s) >= 5 and len(s) <= 20

def gen_large_mpzs(count=default_count):
  return gen_mpzs(mindigits=21,maxdigits=100,count=count)

def is_large_mpz(s):
  return len(s) >= 21

def gen_mpz_spread(count=default_count):
  return gen_small_mpzs(count) + gen_medium_mpzs(count) + gen_large_mpzs(count)

def gen_mpz_args(count=default_count):
    return mpz_std_list + gen_mpz_spread(count)

def gen_mpq_args(count=4):
  nums = zero_one + gen_mpz_spread(count)
  dens = ["1"]    + gen_mpz_spread(count)
  return [n+"/"+d for n in nums for d in dens if int(d) != 0]

def gen_si_args():
  return si_std_list + gen_sis()

def gen_ui_args():
  return ui_std_list + gen_uis()

def gen_list_for_type(t, is_write_only):
  if (t == gmpapi.mpz_t or t == gmpapi.mpq_t) and is_write_only:
    return ["0"]
  elif t == gmpapi.mpz_t:
      return gen_mpz_args()
  elif t == gmpapi.ilong:
    return gen_si_args()
  elif t == gmpapi.ulong:
    return gen_ui_args()
  elif t == gmpapi.mpq_t:
    return gen_mpq_args()
  else:
    raise RuntimeError("Unknown type: {}".format(t))

def gen_args(api):
  if api.custom_test or api.name in custom:
    return custom[api.name](api)
  types = api.params
  if len(types) == 1:
    return [[a] for a in gen_list_for_type(types[0], api.is_write_only(0))]
  elif len(types) == 2:
    t1 = gen_list_for_type(types[0], api.is_write_only(0))
    t2 = gen_list_for_type(types[1], api.is_write_only(1))
    return [(a,b) for a in t1 for b in t2]
  elif len(types) == 3:
    t1 = gen_list_for_type(types[0], api.is_write_only(0))
    t2 = gen_list_for_type(types[1], api.is_write_only(1))
    t3 = gen_list_for_type(types[2], api.is_write_only(2))
    return [(a,b,c) for a in t1 for b in t2 for c in t3]
  elif len(types) == 4:
    t1 = gen_list_for_type(types[0], api.is_write_only(0))
    t2 = gen_list_for_type(types[1], api.is_write_only(1))
    t3 = gen_list_for_type(types[2], api.is_write_only(2))
    t4 = gen_list_for_type(types[3], api.is_write_only(3))
    return [(a,b,c,d) for a in t1 for b in t2 for c in t3 for d in t4]
  else:
    raise RuntimeError("Too many args: {}".format(len(types)))

###################################################################
#
# Fixup and massage random data for better test coverage
#
###################################################################
def mul_mpzs(a, b):
  return str(int(a) * int(b))

def mpz_divexact_data(args):
  # set n = n * d
  divisible = mul_mpzs(args[1],(args[2]))
  return [(args[0], divisible, args[2])]

def mpz_divisible_p_data(args):
  (n,d) = get_div_data(args[0], args[1],rate=1.0)
  return [(n,d), (args[0], args[1])]

def mpz_div3_data(args):
  q     = args[0]
  (n,d) = get_div_data(args[1], args[2],rate=1.0)
  return [(q,n,d), (q,args[1], args[2])]

def mpz_pow_data(args, alwaysallowbase1=True):
  base = int(args[1])
  exp  = int(args[2])
  # allow special numbers
  if base == 0 or exp == 0 or exp == 1:
    return [args]
  if base == 1 and alwaysallowbase1:
    return [args]

  # disallow too big numbers
  if base > 1000 or base < -1000:
    base = gen_random_mpz(maxdigits=3)
  if exp > 1000:
    exp = gen_random_mpz(maxdigits=3, allowneg=False)

  return [(args[0], str(base), str(exp))]

def mpz_mul_2exp_data(args):
  return mpz_pow_data(args, alwaysallowbase1=False)

def mpz_gcd_data(args):
  r = args[0]
  a = args[1]
  b = args[2]
  s_ = gen_small_mpzs(1)[0]
  m_ = gen_medium_mpzs(1)[0]
  l_ = gen_large_mpzs(1)[0]

  return [ (r, a, b),
           (r, mul_mpzs(a, b),  b),
           (r, mul_mpzs(a, s_), mul_mpzs(b, s_)),
           (r, mul_mpzs(a, m_), mul_mpzs(b, m_)),
           (r, mul_mpzs(a, l_), mul_mpzs(b, l_)),]

def mpz_export_data(api):
  rop    = ["0"]
  countp = ["0"]
  order  = ["-1","1"]
  size   = ["1","2","4","8"]
  endian = ["0"]
  nails  = ["0"]
  ops    = gen_mpz_args(1000) + gen_mpzs(count=100,mindigits=100,maxdigits=1000)

  args = []
  for r in rop:
    for c in countp:
      for o in order:
        for s in size:
          for e in endian:
            for n in nails:
              for op in ops:
                args.append((r,c,o,s,e,n,op))
  return args

def mpz_sizeinbase_data(api):
  bases = list(map(str, range(2, 37)))
  ops    = gen_mpz_args(1000) + gen_mpzs(count=1000,mindigits=100,maxdigits=2000)
  return [(op, b) for op in ops for b in bases]

def get_str_data(ty):
  bases = list(range(2,37)) + list(range(-2, -37, -1))
  bases = list(map(str, bases))
  if ty == gmpapi.mpz_t:
    ops = gen_mpz_args(1000)
  elif ty == gmpapi.mpq_t:
    ops = gen_mpq_args(20)
  else:
    raise RuntimeError("Unsupported get_str type: "+str(ty))
  return [("NULL", b, op) for b in bases for op in ops]

def mpz_get_str_data(api):
  return get_str_data(gmpapi.mpz_t)

def mpq_get_str_data(api):
  return get_str_data(gmpapi.mpq_t)

def mpq_set_str_data(api):
  args = gen_mpq_args(20) + gen_mpz_args()
  # zero does not match results exactly because the
  # results are not canonicalized first. We choose to
  # exclude zero from test results. The other option is
  # to canonicalize the results after parsing the strings.
  # Instead we exclude zero so that we can independently
  # test correctness of set_str and canonicalization
  nonzero = []
  for arg in args:
    if "/" in arg:
      pos = arg.find("/")
      if int(arg[:pos]) != 0:
        nonzero.append(arg)
    elif int(arg) != 0:
        nonzero.append(arg)

  return [("0", q, "10") for q in nonzero]

def get_div_data(n, d, rate=0.2):
  """Generate some inputs that are perfectly divisible"""
  if random.random() < rate:
    n = mul_mpzs(n, d)
  return (n,d)

def allow(name, args):
  if name not in blacklists:
    return True
  filters = blacklists[name]
  for (pos, disallow) in filters:
    if args[pos] in disallow:
      return False
  return True

def fixup_args(name, args):
  if name not in fixups:
    return [args]
  return fixups[name](args)


# list of values to be excluded for various api calls
# list format is (pos, [list of values to exclude])
blacklists = {
  "mpz_cdiv_q"      : [(2, ["0", "-0"])],
  "mpz_fdiv_q"      : [(2, ["0", "-0"])],
  "mpz_fdiv_r"      : [(2, ["0", "-0"])],
  "mpz_tdiv_q"      : [(2, ["0", "-0"])],
  "mpz_fdiv_q_ui"   : [(2, ["0", "-0"])],
  "mpz_divexact"    : [(2, ["0", "-0"])],
  "mpz_divisible_p" : [(1, ["0", "-0"])],
  "mpz_divexact_ui" : [(2, ["0", "-0"])],
  "mpq_set_ui"      : [(2, ["0", "-0"])],
}

fixups = {
  "mpz_divexact"    : mpz_divexact_data,
  "mpz_divisible_p" : mpz_divisible_p_data,
  "mpz_cdiv_q"      : mpz_div3_data,
  "mpz_fdiv_q"      : mpz_div3_data,
  "mpz_fdiv_r"      : mpz_div3_data,
  "mpz_tdiv_q"      : mpz_div3_data,
  "mpz_fdiv_q_ui"   : mpz_div3_data,
  "mpz_divexact_ui" : mpz_divexact_data,
  "mpz_pow_ui"      : mpz_pow_data,
  "mpz_gcd"         : mpz_gcd_data,
  "mpz_lcm"         : mpz_gcd_data,
  "mpz_mul_2exp"    : mpz_mul_2exp_data,
}

custom = {
  "mpz_export" : mpz_export_data,
  "mpz_import" : mpz_export_data,
  "mpz_sizeinbase" : mpz_sizeinbase_data,
  "mpz_get_str" : mpz_get_str_data,
  "mpq_set_str" : mpq_set_str_data,
  "mpq_get_str" : mpq_get_str_data,
}

if __name__ == "__main__":
  #apis = [gmpapi.get_api("mpq_set_str"),]
  apis = gmpapi.apis
  for api in apis:
    tests = gen_args(api)
    for args in tests:
      expanded_args = fixup_args(api.name, args)
      for args in expanded_args:
        if allow(api.name, args):
          print("{}|{}".format(api.name, ",".join(args)))

