
import itertools

from brian2.codegen.cpp_prefs import C99Check
from brian2.devices.device import all_devices
from brian2.utils.stringtools import word_substitute, deindent, indent
from brian2.parsing.rendering import NodeRenderer
from brian2.parsing.bast import brian_dtype_from_dtype
from brian2.core.functions import DEFAULT_FUNCTIONS, Function
from brian2.core.variables import (Constant, AuxiliaryVariable,
                                   get_dtype_str, Variable, Subexpression)

from .base import CodeGenerator


__all__ = ['CythonCodeGenerator']


data_type_conversion_table = [
    # canonical         C++            Numpy
    ('float32',        'float',       'float32'),
    ('float64',        'double',      'float64'),
    ('int32',          'int32_t',     'int32'),
    ('int64',          'int64_t',     'int64'),
    ('bool',           'bool',        'bool'),
    ('uint8',          'char',        'uint8'),
    ('uint64',         'uint64_t',    'uint64'),
    ]

cpp_dtype = dict((canonical, cpp) for canonical, cpp, np in data_type_conversion_table)
numpy_dtype = dict((canonical, np) for canonical, cpp, np in data_type_conversion_table)

def get_cpp_dtype(obj):
    return cpp_dtype[get_dtype_str(obj)]

def get_numpy_dtype(obj):
    return numpy_dtype[get_dtype_str(obj)]


class CythonNodeRenderer(NodeRenderer):
    def render_NameConstant(self, node):
        return {True: '1',
                False: '0'}.get(node.value, node.value)

    def render_Name(self, node):
        return {'True': '1',
                'False': '0'}.get(node.id, node.id)

    def render_BinOp(self, node):
        if node.op.__class__.__name__=='Mod':
            left = self.render_node(node.left)
            right = self.render_node(node.right)
            return f'((({left})%({right}))+({right}))%({right})'
        else:
            return super(CythonNodeRenderer, self).render_BinOp(node)


class CythonCodeGenerator(CodeGenerator):
    """
    Cython code generator
    """

    class_name = 'cython'

    def __init__(self, *args, **kwds):
        self.temporary_vars = set()
        super(CythonCodeGenerator, self).__init__(*args, **kwds)

    def translate_expression(self, expr):
        expr = word_substitute(expr, self.func_name_replacements)
        return CythonNodeRenderer(auto_vectorise=self.auto_vectorise).render_expr(expr, self.variables).strip()

    def translate_statement(self, statement):
        var, op, expr, comment = (statement.var, statement.op,
                                  statement.expr, statement.comment)
        if op == ':=': # make no distinction in Cython (declaration are done elsewhere)
            op = '='
        # For Cython we replace complex expressions involving boolean variables into a sequence of
        # if/then expressions with simpler expressions. This is provided by the optimise_statements
        # function.
        if (statement.used_boolean_variables is not None and len(statement.used_boolean_variables)
                # todo: improve dtype analysis so that this isn't necessary
                and brian_dtype_from_dtype(statement.dtype)=='float'):
            used_boolvars = statement.used_boolean_variables
            bool_simp = statement.boolean_simplified_expressions
            codelines = []
            firstline = True
            # bool assigns is a sequence of (var, value) pairs giving the conditions under
            # which the simplified expression simp_expr holds
            for bool_assigns, simp_expr in bool_simp.items():
                # generate a boolean expression like ``var1 and var2 and not var3``
                atomics = []
                for boolvar, boolval in bool_assigns:
                    if boolval:
                        atomics.append(boolvar)
                    else:
                        atomics.append(f"not {boolvar}")
                # use if/else/elif correctly
                if firstline:
                    line = f"if {' and '.join(atomics)}:"
                else:
                    if len(used_boolvars)>1:
                        line = f"elif {' and '.join(atomics)}:"
                    else:
                        line = 'else:'
                line += '\n    '
                line += f"{var} {op} {self.translate_expression(simp_expr)}"
                codelines.append(line)
                firstline = False
            code = '\n'.join(codelines)
        else:
            code = f"{var} {op} {self.translate_expression(expr)}"
        if len(comment):
            code += f" # {comment}"
        return code

    def translate_one_statement_sequence(self, statements, scalar=False):
        # Note that we do not call this function from
        # `translate_statement_sequence` (which has been overwritten)
        # It is nevertheless implemented, so that it can be called explicitly
        # (e.g. from the GSL code generation)
        read, write, indices, conditional_write_vars = self.arrays_helper(statements)
        lines = []
        # index and read arrays (index arrays first)
        lines += self.translate_to_read_arrays(read, indices)
        # the actual code
        lines += self.translate_to_statements(statements, conditional_write_vars)
        # write arrays
        lines += self.translate_to_write_arrays(write)

        return lines

    def translate_to_read_arrays(self, read, indices):
        lines = []
        for varname in itertools.chain(indices, read):
            var = self.variables[varname]
            index = self.variable_indices[varname]
            arrayname = self.get_array_name(var)
            line = f'{varname} = {arrayname}[{index}]'
            lines.append(line)
        return lines

    def translate_to_statements(self, statements, conditional_write_vars):
        lines = []
        for stmt in statements:
            if stmt.op == ':=' and not stmt.var in self.variables:
                self.temporary_vars.add((stmt.var, stmt.dtype))
            line = self.translate_statement(stmt)
            if stmt.var in conditional_write_vars:
                subs = {}
                condvar = conditional_write_vars[stmt.var]
                lines.append(f'if {condvar}:')
                lines.append(indent(line))
            else:
                lines.append(line)
        return lines

    def translate_to_write_arrays(self, write):
        lines = []
        for varname in write:
            index_var = self.variable_indices[varname]
            var = self.variables[varname]
            line = f"{self.get_array_name(var, self.variables)}[{index_var}] = {varname}"
            lines.append(line)
        return lines

    def translate_statement_sequence(self, sc_statements, ve_statements):
        # This function is overwritten, since we do not want to completely
        # separate the code generation for scalar and vector code

        assert set(sc_statements.keys()) == set(ve_statements.keys())

        sc_code = {}
        ve_code = {}

        for block_name in sc_statements:
            sc_block = sc_statements[block_name]
            ve_block = ve_statements[block_name]
            (sc_read, sc_write,
             sc_indices, sc_cond_write) = self.arrays_helper(sc_block)
            (ve_read, ve_write,
             ve_indices, ve_cond_write) = self.arrays_helper(ve_block)
            # We want to read all scalar variables that are needed in the
            # vector code already in the scalar code, if they are not written
            for varname in set(ve_read):
                var = self.variables[varname]
                if var.scalar and varname not in ve_write:
                    sc_read.add(varname)
                    ve_read.remove(varname)

            for (code, stmts, read, write, indices,
                 cond_write) in [(sc_code, sc_block, sc_read, sc_write,
                                  sc_indices, sc_cond_write),
                                 (ve_code, ve_block, ve_read, ve_write,
                                  ve_indices, ve_cond_write)]:
                lines = []
                # index and read arrays (index arrays first)
                lines += self.translate_to_read_arrays(read, indices)
                # the actual code
                lines += self.translate_to_statements(stmts, cond_write)
                # write arrays
                lines += self.translate_to_write_arrays(write)
                code[block_name] = '\n'.join(lines)

        kwds = self.determine_keywords()
        return sc_code, ve_code, kwds

    def _add_user_function(self, varname, var, added):
        user_functions = []
        load_namespace = []
        support_code = []
        impl = var.implementations[self.codeobj_class]
        if (impl.name, var) in added:
            return  # nothing to do
        else:
            added.add((impl.name, var))
        func_code = impl.get_code(self.owner)
        # Implementation can be None if the function is already
        # available in Cython (possibly under a different name)
        if func_code is not None:
            if isinstance(func_code, str):
                # Function is provided as Cython code
                # To make namespace variables available to functions, we
                # create global variables and assign to them in the main
                # code
                user_functions.append((varname, var))
                func_namespace = impl.get_namespace(self.owner) or {}
                for ns_key, ns_value in func_namespace.items():
                    load_namespace.append(
                        f'# namespace for function {varname}')
                    if hasattr(ns_value, 'dtype'):
                        if ns_value.shape == ():
                            raise NotImplementedError((
                            'Directly replace scalar values in the function '
                            'instead of providing them via the namespace'))
                        newlines = [
                            "global _namespace{var_name}",
                            "global _namespace_num{var_name}",
                            "cdef _numpy.ndarray[{cpp_dtype}, ndim=1, mode='c'] _buf_{var_name} = _namespace['{var_name}']",
                            "_namespace{var_name} = <{cpp_dtype} *> _buf_{var_name}.data",
                            "_namespace_num{var_name} = len(_namespace['{var_name}'])"
                        ]
                        support_code.append(
                            f"cdef {get_cpp_dtype(ns_value.dtype)} *_namespace{ns_key}")

                    else:  # e.g. a function
                        newlines = [
                            "_namespace{var_name} = namespace['{var_name}']"
                        ]
                    for line in newlines:
                        load_namespace.append(
                            line.format(cpp_dtype=get_cpp_dtype(ns_value.dtype),
                                        numpy_dtype=get_numpy_dtype(
                                            ns_value.dtype),
                                        var_name=ns_key))
                # Rename references to any dependencies if necessary
                for dep_name, dep in impl.dependencies.items():
                    dep_impl = dep.implementations[self.codeobj_class]
                    dep_impl_name = dep_impl.name
                    if dep_impl_name is None:
                        dep_impl_name = dep.pyfunc.__name__
                    if dep_name != dep_impl_name:
                        func_code = word_substitute(func_code,
                                                    {dep_name: dep_impl_name})
                support_code.append(deindent(func_code))
            elif callable(func_code):
                self.variables[varname] = func_code
                line = f'{varname} = _namespace["{varname}"]'
                load_namespace.append(line)
            else:
                raise TypeError(f"Provided function implementation for function "
                                f"'{varname}' is neither a string nor callable (is "
                                f"type {type(func_code)} instead).")

        dep_support_code = []
        dep_load_namespace = []
        dep_user_functions = []
        if impl.dependencies is not None:
            for dep_name, dep in impl.dependencies.items():
                if dep_name not in self.variables:
                    self.variables[dep_name] = dep
                    user_func = self._add_user_function(dep_name, dep, added)
                    if user_func is not None:
                        sc, ln, uf = user_func
                        dep_support_code.extend(sc)
                        dep_load_namespace.extend(ln)
                        dep_user_functions.extend(uf)

        return (support_code + dep_support_code,
                dep_load_namespace + load_namespace,
                dep_user_functions + user_functions)

    def determine_keywords(self):
        from brian2.devices.device import get_device
        device = get_device()
        # load variables from namespace
        load_namespace = []
        support_code = []
        handled_pointers = set()
        user_functions = []
        added = set()
        for varname, var in sorted(self.variables.items()):
            if isinstance(var, Variable) and not isinstance(var, (Subexpression, AuxiliaryVariable)):
                load_namespace.append(f'_var_{varname} = _namespace["_var_{varname}"]')
            if isinstance(var, AuxiliaryVariable):
                line = f"cdef {get_cpp_dtype(var.dtype)} {varname}"
                load_namespace.append(line)
            elif isinstance(var, Subexpression):
                dtype = get_cpp_dtype(var.dtype)
                line = f"cdef {dtype} {varname}"
                load_namespace.append(line)
            elif isinstance(var, Constant):
                dtype_name = get_cpp_dtype(var.value)
                line = f'cdef {dtype_name} {varname} = _namespace["{varname}"]'
                load_namespace.append(line)
            elif isinstance(var, Variable):
                if var.dynamic:
                    pointer_name = self.get_array_name(var, False)
                    load_namespace.append(f'{pointer_name} = _namespace["{pointer_name}"]')

                # This is the "true" array name, not the restricted pointer.
                array_name = device.get_array_name(var)
                pointer_name = self.get_array_name(var)
                if pointer_name in handled_pointers:
                    continue
                if getattr(var, 'ndim', 1) > 1:
                    continue  # multidimensional (dynamic) arrays have to be treated differently
                if get_dtype_str(var.dtype) == 'bool':
                    newlines = ["cdef _numpy.ndarray[char, ndim=1, mode='c', cast=True] _buf_{array_name} = _namespace['{array_name}']",
                                "cdef {cpp_dtype} * {array_name} = <{cpp_dtype} *> _buf_{array_name}.data"]
                else:
                    newlines = ["cdef _numpy.ndarray[{cpp_dtype}, ndim=1, mode='c'] _buf_{array_name} = _namespace['{array_name}']",
                                "cdef {cpp_dtype} * {array_name} = <{cpp_dtype} *> _buf_{array_name}.data"]

                if not var.scalar:
                    newlines += ["cdef size_t _num{array_name} = len(_namespace['{array_name}'])"]

                if var.scalar and var.constant:
                    newlines += ['cdef {cpp_dtype} {varname} = _namespace["{varname}"]']
                else:
                    newlines += ["cdef {cpp_dtype} {varname}"]

                for line in newlines:
                    line = line.format(cpp_dtype=get_cpp_dtype(var.dtype),
                                       numpy_dtype=get_numpy_dtype(var.dtype),
                                       pointer_name=pointer_name,
                                       array_name=array_name,
                                       varname=varname
                                       )
                    load_namespace.append(line)
                handled_pointers.add(pointer_name)

            elif isinstance(var, Function):
                user_func = self._add_user_function(varname, var, added)
                if user_func is not None:
                    sc, ln, uf = user_func
                    support_code.extend(sc)
                    load_namespace.extend(ln)
                    user_functions.extend(uf)
            else:
                # fallback to Python object
                load_namespace.append(f'{varname} = _namespace["{varname}"]')

        for varname, dtype in sorted(self.temporary_vars):
            cpp_dtype = get_cpp_dtype(dtype)
            line = f"cdef {cpp_dtype} {varname}"
            load_namespace.append(line)

        return {'load_namespace': '\n'.join(load_namespace),
                'support_code_lines': support_code}

###############################################################################
# Implement functions
################################################################################
# Functions that exist under the same name in C++
for func in ['sin', 'cos', 'tan', 'sinh', 'cosh', 'tanh', 'exp', 'log',
             'log10', 'sqrt', 'ceil', 'floor', 'abs']:
    DEFAULT_FUNCTIONS[func].implementations.add_implementation(CythonCodeGenerator,
                                                               code=None)
DEFAULT_FUNCTIONS['expm1'].implementations.add_implementation(CythonCodeGenerator,
                                                              code=None,
                                                              availability_check=C99Check('expm1'))
DEFAULT_FUNCTIONS['log1p'].implementations.add_implementation(CythonCodeGenerator,
                                                              code=None,
                                                              availability_check=C99Check('log1p'))

# Functions that need a name translation
for func, func_cpp in [('arcsin', 'asin'), ('arccos', 'acos'), ('arctan', 'atan'),
                       ('int', 'int_')  # from stdint_compat.h
                       ]:
    DEFAULT_FUNCTIONS[func].implementations.add_implementation(CythonCodeGenerator,
                                                               code=None,
                                                               name=func_cpp)

exprel_code = """
cdef inline double _exprel(double x) nogil:
    if fabs(x) < 1e-16:
        return 1.0
    elif x > 717:  # near log(DBL_MAX)
        return NPY_INFINITY
    else:
        return expm1(x) / x
"""
DEFAULT_FUNCTIONS['exprel'].implementations.add_implementation(CythonCodeGenerator,
                                                               code=exprel_code,
                                                               name='_exprel',
                                                               availability_check=C99Check('exprel'))
_BUFFER_SIZE = 20000

rand_code = """
cdef double _rand(int _idx):
    cdef double **buffer_pointer = <double**>_namespace_rand_buffer
    cdef double *buffer = buffer_pointer[0]
    cdef _numpy.ndarray _new_rand

    if(_namespace_rand_buffer_index[0] == 0):
        if buffer != NULL:
            free(buffer)
        _new_rand = _numpy.random.rand(_BUFFER_SIZE)
        buffer = <double *>_numpy.PyArray_DATA(_new_rand)
        PyArray_CLEARFLAGS(<_numpy.PyArrayObject*>_new_rand, _numpy.NPY_OWNDATA)
        buffer_pointer[0] = buffer

    cdef double val = buffer[_namespace_rand_buffer_index[0]]
    _namespace_rand_buffer_index[0] += 1
    if _namespace_rand_buffer_index[0] == _BUFFER_SIZE:
        _namespace_rand_buffer_index[0] = 0
    return val
""".replace('_BUFFER_SIZE', str(_BUFFER_SIZE))

randn_code = rand_code.replace('rand', 'randn').replace('randnom', 'random')

poisson_code = """
cdef double _loggam(double x):
  cdef double x0, x2, xp, gl, gl0
  cdef int32_t k, n

  cdef double a[10]
  a[:] = [8.333333333333333e-02, -2.777777777777778e-03,
          7.936507936507937e-04, -5.952380952380952e-04,
          8.417508417508418e-04, -1.917526917526918e-03,
          6.410256410256410e-03, -2.955065359477124e-02,
          1.796443723688307e-01, -1.39243221690590e+00]
  x0 = x
  n = 0
  if (x == 1.0) or (x == 2.0):
    return 0.0
  elif x <= 7.0:
    n = <int32_t>(7 - x)
    x0 = x + n
  x2 = 1.0 / (x0 * x0)
  xp = 2 * M_PI
  gl0 = a[9]
  for k in range(8, -1, -1):
    gl0 *= x2
    gl0 += a[k]
  gl = gl0 / x0 + 0.5 * log(xp) + (x0 - 0.5) * log(x0) - x0
  if x <= 7.0:
    for k in range(1, n+1):
      gl -= log(x0 - 1.0)
      x0 -= 1.0
  return gl


cdef int32_t _poisson_mult(double lam, int _vectorisation_idx):
  cdef int32_t X
  cdef double prod, U, enlam

  enlam = exp(-lam)
  X = 0
  prod = 1.0
  while True:
    U = _rand(_vectorisation_idx)
    prod *= U
    if (prod > enlam):
      X += 1
    else:
      return X

cdef int32_t _poisson_ptrs(double lam, int _vectorisation_idx):
  cdef int32_t k
  cdef double U, V, slam, loglam, a, b, invalpha, vr, us

  slam = sqrt(lam)
  loglam = log(lam)
  b = 0.931 + 2.53 * slam
  a = -0.059 + 0.02483 * b
  invalpha = 1.1239 + 1.1328 / (b - 3.4)
  vr = 0.9277 - 3.6224 / (b - 2)

  while True:
    U = _rand(_vectorisation_idx) - 0.5
    V = _rand(_vectorisation_idx)
    us = 0.5 - abs(U)
    k = <int32_t>floor((2 * a / us + b) * U + lam + 0.43)
    if (us >= 0.07) and (V <= vr):
      return k
    if ((k < 0) or ((us < 0.013) and (V > us))):
      continue
    if ((log(V) + log(invalpha) - log(a / (us * us) + b)) <=
        (-lam + k * loglam - _loggam(k + 1))):
      return k

cdef int32_t _poisson(double lam, int32_t _idx):
  if lam >= 10:
    return _poisson_ptrs(lam, _idx)
  elif lam == 0:
    return 0
  else:
    return _poisson_mult(lam, _idx)
"""

device = all_devices['runtime']
DEFAULT_FUNCTIONS['rand'].implementations.add_implementation(CythonCodeGenerator,
                                                             code=rand_code,
                                                             name='_rand',
                                                             namespace={'_rand_buffer': device.rand_buffer,
                                                                        '_rand_buffer_index': device.rand_buffer_index})

DEFAULT_FUNCTIONS['randn'].implementations.add_implementation(CythonCodeGenerator,
                                                              code=randn_code,
                                                              name='_randn',
                                                              namespace={
                                                                  '_randn_buffer': device.randn_buffer,
                                                                  '_randn_buffer_index': device.randn_buffer_index})
DEFAULT_FUNCTIONS['poisson'].implementations.add_implementation(CythonCodeGenerator,
                                                                code=poisson_code,
                                                                name='_poisson',
                                                                dependencies={'_rand': DEFAULT_FUNCTIONS['rand']})

sign_code = """
ctypedef fused _to_sign:
    char
    short
    int
    long
    float
    double

cdef int _sign(_to_sign x):
    return (0 < x) - (x < 0)
"""
DEFAULT_FUNCTIONS['sign'].implementations.add_implementation(CythonCodeGenerator,
                                                             code=sign_code,
                                                             name='_sign')

clip_code = """
ctypedef fused _to_clip:
    char
    short
    int
    long
    float
    double

cdef _to_clip _clip(_to_clip x, double low, double high):
    if x < low:
        return <_to_clip?>low
    if x > high:
        return <_to_clip?>high
    return x
"""
DEFAULT_FUNCTIONS['clip'].implementations.add_implementation(CythonCodeGenerator,
                                                             code=clip_code,
                                                             name='_clip')

timestep_code = """
cdef int64_t _timestep(double t, double dt):
    return <int64_t>((t + 1e-3*dt)/dt)
"""
DEFAULT_FUNCTIONS['timestep'].implementations.add_implementation(CythonCodeGenerator,
                                                                 code=timestep_code,
                                                                 name='_timestep')
