import glob
import os
import re
import sys
import numpy as np

my_path = os.path.dirname(os.path.realpath(__file__))
cppy_files = glob.glob(my_path + '/cppy_source/*.cppy')

dtype_map = {
    'bool': 'bool',
    'float32': 'float',
    'float64': 'double',
    'int32': 'int',
    'int64': 'long',
    'uint32': 'unsigned int',
    'uint64': 'unsigned long',
    'complex64': 'std::complex<float>',
    'complex128': 'std::complex<double>',
}

def get_xtype(arr, xtensor=False):
    if isinstance(arr, (int)):
        return "int"
    elif isinstance(arr, float):
        return "double"

    if xtensor:
        s = "xtensor<" + dtype_map[str(arr.dtype)] + ", " + str(arr.ndim) + ">"
    else:
        s = "xarray<" + dtype_map[str(arr.dtype)] + ">"
    return s

def get_cpp_initlist(arr, name):
    if isinstance(arr, (int, str, float)):
        return get_xtype(arr) + " " + name + " = " + str(arr) + ';'
    name = get_xtype(arr) + " " + name
    s = np.array2string(arr, separator=',', precision=16)
    s = s.replace('[', '{')
    s = s.replace(']', '}')
    s = s.replace('j', 'i')
    s += ';'
    s = s.replace("\n", "\n" + " " * (len(name) + 3))
    s = name + " = " + s
    return s



def translate_file(contents, f):
	current_vars = {}

	matches = re.findall(r"\/\*py.*?\*\/", contents, re.MULTILINE | re.DOTALL)

	def exec_comment(txt, upper_level=False):
		lines = txt.split('\n')
		if upper_level:
			txt = '\n'.join(["import numpy as np"] + [x.strip() for x in lines[1:-1]])
		locals_before = list(locals().keys())
		exec(txt, globals(), current_vars)
		current_vars.update(
			{x: val for x, val in locals().items() if x not in locals_before}
		)

	result_file = ""

	idx = 0
	lidx = 0
	for line in contents.split('\n'):
		if lidx == 8:
			f = os.path.split(f)[1]
			result_file += "// This file is generated from test/files/cppy_source/{} by preprocess.py!\n".format(f)
			result_file += "// Warning: This file should not be modified directly! " \
						   "Instead, modify the `*.cppy` file.\n\n"

		lstrip = line.lstrip()
		if lstrip.startswith("/*py"):
			exec_comment(matches[idx], True)
			idx += 1
		if lstrip.startswith("// py_"):
			indent_n = len(line) - len(lstrip)
			if '=' in lstrip:
				exec_comment(lstrip[6:])
				var = line.strip()[6:lstrip.index('=')].strip()
			else:
				var = line.strip()[6:]
			indent = line[:indent_n]
			init_list = get_cpp_initlist(current_vars[var], 'py_' + var)
			init_list = '\n'.join([indent + x for x in init_list.split('\n')])
			result_file += line + '\n'
			result_file += init_list + '\n'
		else:
			result_file += line + '\n'
		lidx += 1
	return result_file

print("::: PREPROCESSING :::\n")

for f in cppy_files:
	print(" - PROCESSING {}".format(f))

	global current_vars
	current_vars = {}  # reset

	with open(f) as fi:
		contents = fi.read()

	# reset global seed
	np.random.seed(42)
	result = translate_file(contents, f)
	f_result = os.path.split(f)[1]
	with open(my_path + "/../" + f_result[:-1], 'w+') as fo:
		fo.write(result)
	print("::: DONE :::")
	# print(result)