"""
Functions depending on multiple arrays.     GET RID OF THIS FILE
"""


import functools

# from . import gpu_core

CACHE = {}

# This is the implementation of CUDA's atomicAdd in OpenCL
OPENCL_ATOMIC_ADD = """
#pragma OPENCL EXTENSION cl_khr_int64_base_atomics: enable

void atomicAdd(volatile __global double *addr, double val)
{
  union {
    long u;
    double f;
  } next, expected, current;
  current.f = *addr;
  do {
    expected.f = current.f;
    next.f = expected.f + val;
    current.u = atomic_cmpxchg( (volatile __global long *) addr,
				expected.u, next.u);
  } while( current.u != expected.u );
}
"""


def register(function):
  @functools.wraps(function)
  def __wrapper(n):
    if function.__name__ not in CACHE:
      CACHE[function.__name__] = {}
    if n not in CACHE[function.__name__]:
      CACHE[function.__name__][n] = function(n)
    return CACHE[function.__name__][n]

  return __wrapper


def code_for_bin_check(n):
  """
  Generate the code to check whether a value is inside a bin.

  :param n: number of data dimensions.
  :type n: int
  :returns: code.
  :rtype: str
  """
  idata = ",".join(f"GLOBAL_MEM double *data{i}" for i in range(n))
  igaps = ",".join(f"int gap{i}" for i in range(n))
  iedges = ",".join(f"GLOBAL_MEM double *edges{i}" for i in range(n))
  iborders = ",".join(f"GLOBAL_MEM unsigned *border{i}" for i in range(n))

  gap_sizes = """
    int g0 = 1;
    """
  for i in range(1, n):
    gap_sizes += f"""
        int g{i} = gap{i - 1} * g{i - 1};
        """

  conds = " && ".join(
      f"data{i}[idx] >= edges{i}[g{i} * idy] && "
      f"(border{i} ? data{i}[idx] <= edges{i}[g{i} * (idy + 1)] "
      f": data{i}[idx] < edges{i}[g{i} * (idy + 1)])"
      for i in range(n)
  )

  code = f"""
    WITHIN_KERNEL unsigned bin_check( int idx, int idy, {idata}, {igaps}, {iedges}, {iborders} ) {{

       {gap_sizes}

       return ({conds});
    }}
    """
  return code


def build_names(n, name, data_type):
  """
  Build the strings to use for input/output of a variable.

  :param n: number of data dimensions.
  :type n: int
  :param name: name of the variable to use.
  :type name: str
  :param data_type: type of data (double, int, double*, int*, ...).
  :type data_type: str
  :returns: tuple with the names to be used for input and output.
  :rtype: tuple(str, str)
  """
  t = tuple(f"{name}{i}" for i in range(n))
  o = ",".join(t)
  if "*" in data_type:
    i = ",".join(f"GLOBAL_MEM {data_type} {d}" for d in t)
  else:
    i = ",".join(f"{data_type} {d}" for d in t)
  return i, o


@register
def sum_inside_bins(n):
  """
  Create a function to sum occurrences inside bins.

  :param n: number of data parameters.
  :type n: int
  :returns: compiled module.
  :rtype: function
  """
  idata, odata = build_names(n, "in", "double*")
  igaps, ogaps = build_names(n, "gap", "int")
  iedges, oedges = build_names(n, "edges", "double*")
  iborders, oborders = build_names(n, "border", "unsigned*")

  preambulo = "" if gpu_core.BACKEND == gpu_core.CUDA else OPENCL_ATOMIC_ADD

  bin_check = code_for_bin_check(n)

  mod = gpu_core.THREAD.compile(
      f"""
    {preambulo}

    {bin_check}

    KERNEL void sum_inside_bins( GLOBAL_MEM double *out, {idata}, {igaps}, {iedges}, {iborders} )
    {{
       SIZE_T idx = get_global_id(0);
       SIZE_T idy = get_global_id(1);

       if ( bin_check(idx, idy, {odata}, {ogaps}, {oedges}, {oborders}) )
          atomicAdd(&out[idy], 1);
    }}
    """
  )

  return mod.sum_inside_bins


@register
def sum_inside_bins_with_values(n):
  """
  Create a function to sum occurrences inside bins.

  :param n: number of data parameters.
  :type n: int
  :returns: compiled module.
  :rtype: function
  """
  idata, odata = build_names(n, "in", "double*")
  igaps, ogaps = build_names(n, "gap", "int")
  iedges, oedges = build_names(n, "edges", "double*")
  iborders, oborders = build_names(n, "border", "unsigned*")

  preambulo = "" if gpu_core.BACKEND == gpu_core.CUDA else OPENCL_ATOMIC_ADD

  bin_check = code_for_bin_check(n)

  mod = gpu_core.THREAD.compile(
      f"""
    {preambulo}

    {bin_check}

    KERNEL void sum_inside_bins_with_values( GLOBAL_MEM double *out, {idata}, {igaps}, {iedges}, {iborders}, GLOBAL_MEM double *values )
    {{
       SIZE_T idx = get_global_id(0);
       SIZE_T idy = get_global_id(1);

       if ( bin_check(idx, idy, {odata}, {ogaps}, {oedges}, {oborders}) )
          atomicAdd(&out[idy], values[idx]);
    }}
    """
  )

  return mod.sum_inside_bins_with_values
