# Copyright (c) 2022 Graphcore Ltd. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import cppimport

import sys
import json
import os

from typing import Any, Optional, Sequence, Dict
from jax.interpreters import xla
from jax.interpreters import mlir
from jax.lib import xla_client, xla_extension

import jax._src.util as util
from jax._src.lib.mlir import ir
from jax._src.lib.mlir.dialects import mhlo

from .xla_utils import dtype_to_tf_datatype_enum, aval_to_xla_shape, ShapedArray

# Pybind11 extension import (and compilation if necessary).
# Explicit path is more robust to different `pip install` usages.
ext_filename = os.path.abspath(
    os.path.join(os.path.dirname(__file__), "ipu_custom_primitive_impl.cpp")
)
ipu_custom_primitive_impl = cppimport.imp_from_filepath(
    ext_filename, "jax.ipu.primitive.ipu_custom_primitive_impl"
)
# Import pybind11 classes.
PrimitiveMetadata = ipu_custom_primitive_impl.PrimitiveMetadata


def make_custom_primitive_attributes(
    primitive_cls: Any, avals_out: Sequence[ShapedArray],
    opaque_attributes: Optional[str], ipu_gp_filename: Optional[str]
) -> Dict[str, Any]:
  """Build custom primitive attributes dict, following the
  interface defined in the IPU XLA backend.
  Args:
    primitive_cls: Pybind11 C++ primitive class.
    avals_out: Output shaped arrays.
    opaque_attributes: Optional (raw) attributes to pass to the primitive, as a string.
    ipu_gp_filename: (Optional) IPU custom vertex GP filename. If a C++ filename is passed,
      it will be automatically compiled by `popc`. Filename should be relative to
      the directory of the main C++ primitive file or absolute.
  Returns:
    Custom op attributes dict.
  """
  # Extract proper information from the primitive binded class.
  op_module = sys.modules[primitive_cls.__module__]
  assert op_module.__file__ is not None
  op_module_path = os.path.abspath(op_module.__file__)
  # Get absolute IPU vertex gp path, if passed as argument.
  if ipu_gp_filename is not None and not os.path.isabs(ipu_gp_filename):
    ipu_gp_filename = os.path.join(
        os.path.dirname(op_module_path), os.path.basename(ipu_gp_filename)
    )
  # Standard naming generated by the macro `EXPORT_IPU_JAX_PRIMITIVE`
  op_name = f"{primitive_cls.__name__}Export"
  opattributes = {
      "op_name": op_name,
      "library_path": op_module_path,
      "gp_path": ipu_gp_filename,
      "output_types": [dtype_to_tf_datatype_enum(v.dtype) for v in avals_out],
      "output_shapes": [v.shape for v in avals_out],
      "is_user_read_write": False,
      "attributes": opaque_attributes,
      # Gradients attributes: unnecessary for JAX.
      "gradient_size": 0,
      "partial_derivative_index": 0,
      "separate_gradients": False,
  }
  return opattributes


def ipu_mlir_lowering_custom_primitive(
    primitive_cls: Any,
    ctx: mlir.LoweringRuleContext,
    operands: Sequence[ir.Value],
    opaque_attributes: Optional[str] = None,
    ipu_gp_filename: Optional[str] = None
) -> Sequence[ir.Value]:
  """IPU MLIR lowering of custom primitive call.
  This method is generating the proper MHLO `CustomCall` to run a custom
  IPU primitive (i.e. passing the proper attributes & metadata).
  Args:
    primitive_cls: Pybind11 C++ primitive class.
    ctx: MLIR translation context.
    operands: List of MLIR input operands.
    opaque_attributes: Optional (raw) attributes to pass to the C++ primitive, as a string.
    ipu_gp_filename: (Optional) IPU custom vertex GP filename. If a C++ filename is passed,
      it will be automatically compiled by `popc`. Filename should be relative to
      the directory of the main C++ primitive file.
  Returns:
    List of MLIR outputs from the `CustomCall`.
  """
  # Build MLIR tuple return type.
  result_types = util.flatten([mlir.aval_to_ir_types(aval) for aval in ctx.avals_out])
  result_type = ir.TupleType.get_tuple(result_types)
  # Custom op/primitive attributes.
  opattributes = make_custom_primitive_attributes(
      primitive_cls, ctx.avals_out, opaque_attributes, ipu_gp_filename
  )
  # IPU XLA backend custom op static name.
  call_target_name = "UserOp"
  # IPU XLA custom op expecting backend attributes encoded as json.
  backend_config = json.dumps(opattributes)
  has_side_effect = False
  result = mhlo.CustomCallOp(
      [result_type],
      operands,
      call_target_name=ir.StringAttr.get(call_target_name),
      has_side_effect=ir.BoolAttr.get(has_side_effect),
      api_version=mlir.i32_attr(0),  # original CustomCall API.
      called_computations=None,
      backend_config=ir.StringAttr.get(backend_config),
      operand_layouts=None,
      result_layouts=None
  )
  # Unpack the output tuple.
  results = [
      mhlo.GetTupleElementOp(result, mlir.i32_attr(i)).result
      for i in range(len(result_types))
  ]
  return results