#  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
#  See https://llvm.org/LICENSE.txt for license information.
#  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

from npcomp.native.mlir import ir

__all__ = [
  "load_builtin_module",
  "DialectHelper",
]


class DialectHelper(ir.DialectHelper):
  r"""Dialect helper.
  
    >>> c = ir.MLIRContext()
    >>> h = DialectHelper(c)
    >>> m = c.new_module()
    >>> tensor_type = h.tensor_type(h.f32_type)
    >>> h.builder.insert_block_start(m.first_block)
    >>> f = h.func_op("foobar", h.function_type(
    ...   [tensor_type, tensor_type], [tensor_type]), 
    ...   create_entry_block=True)
    >>> uf = h.numpy_ufunc_call_op("numpy.add", tensor_type,
    ...   *f.first_block.args)
    >>> _ = h.return_op(uf.results)
    >>> print(m.to_asm())
    <BLANKLINE>
    <BLANKLINE>
    module {
      func @foobar(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>) -> tensor<*xf32> {
        %0 = numpy.ufunc_call @numpy.add(%arg0, %arg1) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
        return %0 : tensor<*xf32>
      }
    }

  Types:
    >>> t = DialectHelper(ir.MLIRContext())
    >>> t.numpy_any_dtype
    !numpy.any_dtype
    >>> t.tensor_type(t.numpy_any_dtype, [1, 2, 3])
    tensor<1x2x3x!numpy.any_dtype>
    >>> t.tensor_type(t.numpy_any_dtype)
    tensor<*x!numpy.any_dtype>
    >>> t.tensor_type(t.numpy_any_dtype, [-1, 2])
    tensor<?x2x!numpy.any_dtype>
    >>> t.tensor_type(t.f32_type)
    tensor<*xf32>
    >>> t.function_type([t.i32_type], [t.f32_type])
    (i32) -> f32

  """
  def __init__(self, *args, **kwargs):
    super().__init__(*args, **kwargs)
    self.numpy_any_dtype = self.context.parse_type("!numpy.any_dtype")

  def numpy_ufunc_call_op(self, callee_symbol, result_type, *args):
    """Creates a numpy.ufunc_call op."""
    c = self.context
    attrs = c.dictionary_attr({
      "ufunc_ref": c.flat_symbol_ref_attr(callee_symbol)
    })
    return self.op("numpy.ufunc_call", [result_type], args, attrs)

  def numpy_narrow_op(self, result_type, operand):
    """Creates a numpy.narrow op."""
    return self.op("numpy.narrow", [result_type], [operand])


def load_builtin_module(context=None):
  """Loads a module populated with numpy built-ins.

  This is not a long-term solution but overcomes some bootstrapping
  issues.

    >>> m = load_builtin_module()
    >>> op = m.region(0).blocks.front.operations.front
    >>> op.is_registered
    True
    >>> op.name
    'numpy.generic_ufunc'

  Args:
    context: The MLIRContext to use (None to create a new one).
  Returns:
    A ModuleOp.
  """
  if context is None: context = ir.MLIRContext()
  return context.parse_asm(_BUILTIN_MODULE_ASM)


_BUILTIN_MODULE_ASM = r"""
  numpy.generic_ufunc @numpy.add (
    overload(%arg0: i32, %arg1: i32) -> i32 {
      %0 = addi %arg0, %arg1 : i32
      numpy.ufunc_return %0 : i32
    },
    overload(%arg0: f32, %arg1: f32) -> f32 {
      %0 = addf %arg0, %arg1 : f32
      numpy.ufunc_return %0 : f32
    }
  )
  numpy.generic_ufunc @numpy.multiply (
    overload(%arg0: i32, %arg1: i32) -> i32 {
      %0 = muli %arg0, %arg1 : i32
      numpy.ufunc_return %0 : i32
    },
    overload(%arg0: f32, %arg1: f32) -> f32 {
      %0 = mulf %arg0, %arg1 : f32
      numpy.ufunc_return %0 : f32
    }
  )
"""

if __name__ == "__main__":
  import doctest
  doctest.testmod()