import sys
import argparse
import collections

arg_p3 = 'Unit *z, const Unit *x, const Unit *y'
arg_p2 = 'Unit *y, const Unit *x'
arg_p2u = 'Unit *z, const Unit *x, Unit y'
param_u3 = 'z, x, y'
param_u2 = 'y, x'

protoType = {
  ('Unit', arg_p3) : 'u_ppp',
  ('void', arg_p3) : 'void_ppp',
  ('void', arg_p2) : 'void_pp',
  ('Unit', arg_p2u) : 'u_ppu',
}

FuncType = collections.namedtuple('FuncType', 'name, ret, args, cname, prams, N, N64')

# print 'Unit mclb_add1(Unit *z, const Unit *x, const Unit *y);'
FuncType.printProto = lambda self, i: print(f'{self.ret} {self.cname}{i}({self.args});')

FuncType.printProtoSF = lambda self, i, suf: print(f'{self.ret} {self.cname}_{suf}{i}({self.args});')

FuncType.getProtoType = lambda self: protoType[(self.ret, self.args)]

def gen_func(name, ret, args, cname, params, i):
  print(f'{ret} {cname}{i}({args});')

def gen_prototype(out, e):
  (name, ret, args, _, params, N, N64) = e

  print(f'''{e.getProtoType()} get_{name}(size_t n)
{{
#if MCL_BINT_ASM == 1''')
  for i in range(1, N):
    if i == N64 + 1:
      print('#if MCL_SIZEOF_UNIT == 4')
    print(f'\tif (n == {i}) return mclb_{name}{i};')
  print('#endif // MCL_SIZEOF_UNIT == 4')
  print('#else // MCL_FP_BITN_ASM == 1')
  for i in range(1, N):
    if i == N64 + 1:
      print('#if MCL_SIZEOF_UNIT == 4')
    print(f'\tif (n == {i}) return {name}T<{i}>;')
  print('''#endif // MCL_SIZEOF_UNIT == 4
#endif // MCL_BINT_ASM == 1
	// CYBOZU_ASSUME(false);
  return 0;
}''')

def roundup(x, n):
  return (x + n - 1) // n

def gen_disable(N, tbl):
  name1 = 'mulUnit'
  name2 = 'mulUnitAdd'
  print('#if MCL_BINT_ASM_X64 == 1')

  print('extern "C" {')
  for i in range(1, N+1):
    for e in tbl:
      e.printProtoSF(i, 'slow')
      e.printProtoSF(i, 'fast')
  print('}')

  print('#ifdef _WIN32')
  print('static const bool g_adx = g_cpuType & tAVX_BMI2_ADX;')
  for i in range(1, N+1):
    for e in tbl:
      print(f'const {e.getProtoType()} {e.cname}{i} = g_adx ? {e.cname}_fast{i} : {e.cname}_slow{i};')
  print('extern "C" void mclb_enable_fast() {')
  print('}')
  print('#else')
  for i in range(1, N+1):
    for e in tbl:
      print(f'{e.getProtoType()} {e.cname}{i} = {e.cname}_slow{i};')
  print('extern "C" void mclb_enable_fast() {')
  for i in range(1, N+1):
    for e in tbl:
      print(f'\t{e.cname}{i} = {e.cname}_fast{i};')
  print('}')
  print('#endif')

  print('#endif // MCL_BINT_ASM_X64 == 1')

def gen_mul_slow(N):
  print('#if MCL_BINT_ASM_X64 == 1')
  for n in range(1,N+1):
    print(f'''extern "C" void mclb_mul_slow{n}(Unit *z, const Unit *x, const Unit *y)
{{
	z[{n}] = mulUnitT<{n}>(z, x, y[0]);
	u_ppu f = get_mulUnitAdd({n});
	for (size_t i = 1; i < {n}; i++) {{
		z[{n} + i] = f(&z[i], x, y[i]);
	}}
}}''')
  print('#endif // MCL_BINT_ASM_X64 == 1')

def gen_sqr_slow(N):
  print('#if MCL_BINT_ASM_X64 == 1')
  for n in range(1,N+1):
    print(f'''extern "C" void mclb_sqr_slow{n}(Unit *y, const Unit *x)
{{
	mclb_mul_slow{n}(y, x, x);
}}''')
  print('#endif // MCL_BINT_ASM_X64 == 1')

def main():
  parser = argparse.ArgumentParser(description='gen header')
  parser.add_argument('out', type=str)
  parser.add_argument('-max_bit', type=int, default=512+32)
  opt = parser.parse_args()
  if not opt.out in ['switch']:
    print('bad out', opt.out)
    sys.exit(1)
  N = roundup(opt.max_bit, 32)
  N64 = roundup(opt.max_bit, 64)
  addN = 32
  addN64 = 16

  addTbl = list(map(lambda x: FuncType(*x), [
    ('add', 'Unit', arg_p3, 'mclb_add', param_u3, addN, addN64),
    ('sub', 'Unit', arg_p3, 'mclb_sub', param_u3, addN, addN64),
    ('addNF', 'void', arg_p3, 'mclb_addNF', param_u3, addN, addN64),
    ('subNF', 'Unit', arg_p3, 'mclb_subNF', param_u3, addN, addN64),
  ]))

  mulTbl = list(map(lambda x: FuncType(*x), [
    ('mulUnit', 'Unit', arg_p2u, 'mclb_mulUnit', param_u3, N, N64),
    ('mulUnitAdd', 'Unit', arg_p2u, 'mclb_mulUnitAdd', param_u3, N, N64),
    ('mul', 'void', arg_p3, 'mclb_mul', param_u3, N, N64),
    ('sqr', 'void', arg_p2, 'mclb_sqr', param_u2, N, N64),
  ]))

  print('// this code is generated by python3 src/gen_bint_header.py', opt.out)
  print('#if MCL_BINT_ASM == 1')
  print('extern "C" {')
  for i in range(1, addN+1):
    if i == addN64 + 1:
      print('#if MCL_SIZEOF_UNIT == 4')
    for e in addTbl:
      e.printProto(i)
  print('#endif // #if MCL_SIZEOF_UNIT == 4')
  print('#if MCL_BINT_ASM_X64 != 1')
  for i in range(1, N+1):
    if i == N64 + 1:
      print('#if MCL_SIZEOF_UNIT == 4')
    for e in mulTbl:
      e.printProto(i)
  print('#endif // #if MCL_SIZEOF_UNIT == 4')
  print('#endif')
  print('}')
  print('#endif // #if MCL_BINT_ASM == 1')

  gen_disable(N64, mulTbl)
  gen_mul_slow(N64)
  gen_sqr_slow(N64)

  for ft in addTbl + mulTbl:
    gen_prototype(opt.out, ft)

if __name__ == '__main__':
  main()

