# 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
import operator
from abc import abstractmethod
from copy import deepcopy
from functools import cached_property, partialmethod
from typing import Optional, Tuple, Union

import numpy as np

from ..util import infer_mlir_type, mlir_type_to_np_dtype
from numba_cuda_mlir._mlir._mlir_libs._mlir import register_value_caster
from numba_cuda_mlir._mlir.dialects import complex as complex_dialect
from numba_cuda_mlir._mlir.dialects import math as math_dialect
from numba_cuda_mlir._mlir.dialects._arith_enum_gen import (
    _arith_cmpfpredicateattr,
    _arith_cmpipredicateattr,
)
from numba_cuda_mlir._mlir.dialects._ods_common import (
    get_op_result_or_op_results,
    get_op_result_or_value,
)
from numba_cuda_mlir._mlir.dialects.arith import *
from numba_cuda_mlir._mlir.ir import (
    Attribute,
    BF16Type,
    ComplexType,
    Context,
    DenseElementsAttr,
    F16Type,
    F32Type,
    F64Type,
    FloatAttr,
    FloatType,
    IndexType,
    InsertionPoint,
    IntegerType,
    Location,
    OpView,
    Operation,
    ShapedType,
    Type,
    Value,
    register_attribute_builder,
)


def constant(
    value: Union[int, float, bool, np.ndarray],
    type: Optional[Type] = None,
    index: Optional[bool] = None,
    *,
    result: Optional[Type] = None,
    vector: Optional[bool] = False,
    loc: Location = None,
    ip: InsertionPoint = None,
) -> Value:
    """Instantiate arith.constant with value `value`.

    Args:
      value: Python value that determines the value attribute of the
        arith.constant op.
      type: Optional MLIR type that type of the value attribute of the
        arith.constant op; if omitted the type of the value attribute
        will be inferred from the value.
      index: Whether the MLIR type should be an index type; if passed the
        type argument will be ignored.

    Returns:
      ir.OpView instance that corresponds to instantiated arith.constant op.
    """
    if index is not None and index:
        type = IndexType.get()

    if type is not None and result is not None:
        raise ValueError("only type XOR result expected")
    type = type or result
    if type is None:
        type = infer_mlir_type(value, vector=vector)

    assert type is not None

    if isinstance(type, ComplexType):
        value = complex(value)
        return get_op_result_or_op_results(
            complex_dialect.ConstantOp(
                type,
                list(
                    map(
                        lambda x: FloatAttr.get(type.element_type, x),
                        [value.real, value.imag],
                    )
                ),
                loc=loc,
                ip=ip,
            )
        )

    if isinstance(type, FloatType) and not isinstance(value, np.ndarray):
        value = float(value)

    if isinstance(type, ShapedType) and isinstance(value, (int, float, bool)):
        ranked_tensor_type = ShapedType(type)
        value = np.full(
            ranked_tensor_type.shape,
            value,
            dtype=mlir_type_to_np_dtype(ranked_tensor_type.element_type),
        )

    if isinstance(value, np.ndarray):
        value = DenseElementsAttr.get(
            value,
            type=type,
        )

    return get_op_result_or_op_results(ConstantOp(type, value, loc=loc, ip=ip))


def index_cast(
    in_: Value,
    to: Type = None,
    *,
    out: Type = None,
    loc: Location = None,
    ip: InsertionPoint = None,
) -> Value:
    assert bool(to) != bool(out), "either `to` or `out` must be set but not both"
    res_type = out or to
    if res_type is None:
        res_type = IndexType.get()
    return get_op_result_or_op_results(IndexCastOp(res_type, in_, loc=loc, ip=ip))


@register_attribute_builder("Arith_CmpIPredicateAttr", replace=True)
def _arith_CmpIPredicateAttr(predicate: Union[str, Attribute], context: Context):
    predicates = {
        "eq": CmpIPredicate.eq,
        "ne": CmpIPredicate.ne,
        "slt": CmpIPredicate.slt,
        "sle": CmpIPredicate.sle,
        "sgt": CmpIPredicate.sgt,
        "sge": CmpIPredicate.sge,
        "ult": CmpIPredicate.ult,
        "ule": CmpIPredicate.ule,
        "ugt": CmpIPredicate.ugt,
        "uge": CmpIPredicate.uge,
        0: CmpIPredicate.eq,
        1: CmpIPredicate.ne,
        2: CmpIPredicate.slt,
        3: CmpIPredicate.sle,
        4: CmpIPredicate.sgt,
        5: CmpIPredicate.sge,
        6: CmpIPredicate.ult,
        7: CmpIPredicate.ule,
        8: CmpIPredicate.ugt,
        9: CmpIPredicate.uge,
    }
    if isinstance(predicate, Attribute):
        return predicate
    predicate = str(predicate)
    assert predicate in predicates, f"{predicate=} not in predicates"
    return _arith_cmpipredicateattr(predicates[predicate], context)


@register_attribute_builder("Arith_CmpFPredicateAttr", replace=True)
def _arith_CmpFPredicateAttr(predicate: Union[str, Attribute], context: Context):
    predicates = {
        "false": CmpFPredicate.AlwaysFalse,
        # ordered comparison
        # An ordered comparison checks if neither operand is NaN.
        "oeq": CmpFPredicate.OEQ,
        "ogt": CmpFPredicate.OGT,
        "oge": CmpFPredicate.OGE,
        "olt": CmpFPredicate.OLT,
        "ole": CmpFPredicate.OLE,
        "one": CmpFPredicate.ONE,
        # no clue what this one is
        "ord": CmpFPredicate.ORD,
        # unordered comparison
        # Conversely, an unordered comparison checks if either operand is a NaN.
        "ueq": CmpFPredicate.UEQ,
        "ugt": CmpFPredicate.UGT,
        "uge": CmpFPredicate.UGE,
        "ult": CmpFPredicate.ULT,
        "ule": CmpFPredicate.ULE,
        "une": CmpFPredicate.UNE,
        # no clue what this one is
        "uno": CmpFPredicate.UNO,
        # return always true
        "true": CmpFPredicate.AlwaysTrue,
    }
    if isinstance(predicate, Attribute):
        return predicate
    predicate = str(predicate)
    assert predicate in predicates, f"{predicate=} not in predicates"
    return _arith_cmpfpredicateattr(predicates[predicate], context)


def _binary_op(
    lhs: "ArithValue",
    rhs: "ArithValue",
    op: str,
    predicate: str = None,
    signedness: str = None,
    *,
    loc: Location = None,
) -> "ArithValue":
    """Generic for handling infix binary operator dispatch.

    Args:
      lhs: E.g. ScalarValue or TensorValue below.
      rhs: ScalarValue or TensorValue with type matching self.
      op: Binary operator, currently only add, sub, mul
        supported.

    Returns:
      Result of binary operation. This will be a handle to an arith(add|sub|mul) op.
    """
    if (isinstance(lhs, Value) and isinstance(rhs, Value) and lhs.type != rhs.type) or isinstance(
        rhs, (float, int, bool, np.ndarray)
    ):
        rhs = lhs.coerce(rhs)
    elif isinstance(rhs, Value) and isinstance(lhs, (float, int, bool, np.ndarray)):
        lhs = rhs.coerce(lhs)

    assert lhs.type == rhs.type, f"{lhs=} {rhs=} must have the same type."

    assert op in {
        "add",
        "and",
        "or",
        "xor",
        "sub",
        "mul",
        "cmp",
        "truediv",
        "floordiv",
        "mod",
    }

    if op == "cmp":
        assert predicate is not None

    if lhs.fold() and lhs.fold():
        klass = lhs.__class__
        # if both operands are constants (results of an arith.constant op)
        # then both have a literal value (i.e. Python value).
        lhs, rhs = lhs.literal_value, rhs.literal_value
        # if we're folding constants (self._fold = True) then we just carry out
        # the corresponding operation on the literal values; e.g., operator.add.
        # note this is the same as op = operator.__dict__[op].
        if predicate is not None:
            op = predicate
        op = operator.attrgetter(op)(operator)
        return klass(op(lhs, rhs), fold=True)

    if op == "truediv":
        op = "div"
    if op == "mod":
        op = "rem"

    op = op.capitalize()
    if isinstance(lhs.dtype, FloatType):
        if op == "Xor":
            raise TypeError(
                f"unsupported operand type(s) for ^: '{type(lhs).__name__}' and "
                f"'{type(rhs).__name__}'"
            )
        if op == "Floordiv":
            # Python float // float is floor(a / b); arith has no floordivf — match
            # numba_cuda_mlir/lowering/math.py floordiv_cg: divf then math.floor.
            div_op = globals()["DivFOp"]
            div_val = lhs.__class__(div_op(lhs, rhs, loc=loc), dtype=lhs.dtype)
            return lhs.__class__(math_dialect.floor(div_val, loc=loc), dtype=lhs.dtype)
        op += "F"
    elif isinstance(lhs.dtype, (IntegerType, IndexType)):
        if "div" in op.lower() or "rem" in op.lower():
            if not isinstance(lhs.dtype, IndexType) and not lhs.dtype.is_signless:
                raise ValueError(f"{op.lower()}i not supported for {lhs=}")
            if op == "Floordiv":
                op = "FloorDiv"
            op += "S"
        op += "I"
    else:
        raise NotImplementedError(f"Unsupported '{op}' operands: {lhs}, {rhs}")

    # because up above we have from arith import *
    # MLIR ODS name is XOrIOp (capital O), not XorIOp, after "xor" -> capitalize -> "Xor" + "I".
    op_cls = f"{op}Op"
    if op_cls == "XorIOp":
        op_cls = "XOrIOp"
    op = globals()[op_cls]

    if predicate is not None:
        if isinstance(lhs.dtype, FloatType):
            # ordered comparison - see above
            predicate = "o" + predicate
        elif isinstance(lhs.dtype, (IntegerType, IndexType)):
            # eq, ne signs don't matter
            if predicate not in {"eq", "ne"}:
                if signedness is not None:
                    predicate = signedness + predicate
                else:
                    if isinstance(lhs.dtype, IndexType) or lhs.dtype.is_unsigned:
                        predicate = "u" + predicate
                    else:
                        assert lhs.dtype.is_signed or lhs.dtype.is_signless
                        predicate = "s" + predicate
        return lhs.__class__(op(predicate, lhs, rhs, loc=loc), dtype=lhs.dtype)
    else:
        return lhs.__class__(op(lhs, rhs, loc=loc), dtype=lhs.dtype)


def _rbinary_op(
    rhs: "ArithValue",
    lhs: "ArithValue",
    op: str,
    predicate: str = None,
    signedness: str = None,
    *,
    loc: Location = None,
) -> "ArithValue":
    return _binary_op(lhs, rhs, op, predicate, signedness, loc=loc)


def _pow_op(
    lhs: "ArithValue",
    rhs: "ArithValue",
    *,
    loc: Location = None,
) -> "ArithValue":
    if isinstance(rhs, (float, int, bool)):
        rhs = lhs.coerce(rhs)
    elif isinstance(lhs, (float, int, bool)):
        lhs = rhs.coerce(lhs)

    if lhs.fold() and rhs.fold():
        klass = lhs.__class__
        return klass(operator.pow(lhs.literal_value, rhs.literal_value), fold=True)

    if isinstance(lhs.dtype, FloatType) and isinstance(rhs.dtype, FloatType):
        return lhs.__class__(math_dialect.powf(lhs, rhs, loc=loc), dtype=lhs.dtype)
    elif isinstance(lhs.dtype, FloatType) and isinstance(rhs.dtype, IntegerType):
        return lhs.__class__(math_dialect.fpowi(lhs, rhs, loc=loc), dtype=lhs.dtype)
    elif isinstance(lhs.dtype, IntegerType) and isinstance(rhs.dtype, IntegerType):
        assert lhs.type == rhs.type, f"{lhs=} {rhs=} must have the same type."
        from numba_cuda_mlir.mlir.dialect_exts.math import ipowi

        return lhs.__class__(ipowi(lhs, rhs, loc=loc), dtype=lhs.dtype)
    else:
        raise NotImplementedError(f"Unsupported pow operands: {lhs.type}, {rhs.type}")


def _rpow_op(
    rhs: "ArithValue",
    lhs: "ArithValue",
    *,
    loc: Location = None,
) -> "ArithValue":
    return _pow_op(lhs, rhs, loc=loc)


class ArithValue(Value, metaclass=type(Value)):
    """Class for functionality shared by Value subclasses that support
    arithmetic operations.

    Wrapping logic is implemented in `__new__` so that nanobind's own
    instance-creation path (`nb_type.__call__`) is used. That path knows how
    to initialize a Python subclass of a nanobind-bound `Value` correctly;
    calling `cls.__init__` explicitly from a Python metaclass does not.

    The class uses `type(Value)` (nanobind's `nb_type`) as its metaclass.
    This lets nanobind recognise `ArithValue` and its subclasses as the
    Python wrapper for MLIR `Value` objects, so block arguments are
    automatically wrapped, while avoiding the metaclass-subclass check that
    nanobind 2.13 tightened.
    """

    def __new__(cls, arg, *, dtype=None, fold=None, loc=None, ip=None):
        if isinstance(arg, (OpView, Operation, Value)):
            if isinstance(arg, (Operation, OpView)):
                assert len(arg.results) == 1
            val = get_op_result_or_value(arg)
        elif isinstance(arg, (int, float, bool, np.ndarray)):
            if dtype is not None and not isinstance(dtype, Type):
                raise ValueError(f"{dtype=} is expected to be an ir.Type.")
            if fold is not None and not isinstance(fold, bool):
                raise ValueError(f"{fold=} is expected to be a bool.")
            val = constant(deepcopy(arg), dtype, loc=loc, ip=ip)
        else:
            raise NotImplementedError(f"{cls.__name__} doesn't support wrapping {arg}.")
        instance = Value.__new__(cls, val)
        instance._arith_value = val
        return instance

    def __init__(self, arg, *, dtype=None, fold: Optional[bool] = None, loc=None, ip=None):
        super().__init__(self._arith_value)
        del self._arith_value
        self._fold = fold if fold is not None else False

    def is_constant(self) -> bool:
        return isinstance(self.owner, OpView) and isinstance(self.owner, ConstantOp)

    @property
    @abstractmethod
    def literal_value(self):
        pass

    @abstractmethod
    def coerce(self, other) -> Tuple["ArithValue", "ArithValue"]:
        pass

    def __hash__(self):
        return Value(self).__hash__()

    def fold(self) -> bool:
        return self.is_constant() and self._fold

    def __str__(self):
        return f"{self.__class__.__name__}({self.get_name()}, {self.type})"

    def __repr__(self):
        return str(Value(self)).replace("Value", self.__class__.__name__)

    # partialmethod differs from partial in that it also binds the object instance
    # to the first arg (i.e., self)
    __add__ = partialmethod(_binary_op, op="add")
    __sub__ = partialmethod(_binary_op, op="sub")
    __mul__ = partialmethod(_binary_op, op="mul")
    __truediv__ = partialmethod(_binary_op, op="truediv")
    __floordiv__ = partialmethod(_binary_op, op="floordiv")
    __mod__ = partialmethod(_binary_op, op="mod")
    __and__ = partialmethod(_binary_op, op="and")
    __or__ = partialmethod(_binary_op, op="or")
    __xor__ = partialmethod(_binary_op, op="xor")
    __pow__ = partialmethod(_pow_op)

    __radd__ = partialmethod(_rbinary_op, op="add")
    __rsub__ = partialmethod(_rbinary_op, op="sub")
    __rmul__ = partialmethod(_rbinary_op, op="mul")
    __rtruediv__ = partialmethod(_rbinary_op, op="truediv")
    __rfloordiv__ = partialmethod(_rbinary_op, op="floordiv")
    __rmod__ = partialmethod(_rbinary_op, op="mod")
    __rand__ = partialmethod(_rbinary_op, op="and")
    __ror__ = partialmethod(_rbinary_op, op="or")
    __rxor__ = partialmethod(_rbinary_op, op="xor")
    __rpow__ = partialmethod(_rpow_op)

    def __eq__(self, other):
        if not isinstance(other, self.__class__):
            try:
                other = self.__class__(other, dtype=self.type)
            except NotImplementedError as e:
                assert "doesn't support wrapping" in str(e)
                return False
        if self is other:
            return True
        return _binary_op(self, other, op="cmp", predicate="eq")

    def __ne__(self, other):
        if not isinstance(other, self.__class__):
            try:
                other = self.__class__(other, dtype=self.type)
            except NotImplementedError as e:
                assert "doesn't support wrapping" in str(e)
                return True
        if self is other:
            return False
        return _binary_op(self, other, op="cmp", predicate="ne")

    __le__ = partialmethod(_binary_op, op="cmp", predicate="le")
    __lt__ = partialmethod(_binary_op, op="cmp", predicate="lt")
    __ge__ = partialmethod(_binary_op, op="cmp", predicate="ge")
    __gt__ = partialmethod(_binary_op, op="cmp", predicate="gt")

    __rle__ = partialmethod(_rbinary_op, op="cmp", predicate="le")
    __rlt__ = partialmethod(_rbinary_op, op="cmp", predicate="lt")
    __rge__ = partialmethod(_rbinary_op, op="cmp", predicate="ge")
    __rgt__ = partialmethod(_rbinary_op, op="cmp", predicate="gt")

    def _eq(self, other):
        return Value(self) == Value(other)

    def _ne(self, other):
        return Value(self) != Value(other)


class ScalarValue(ArithValue):
    """Value subclass ScalarValue that adds convenience methods
    for getting dtype and (possibly) the stored literal value.

    Note, order matters in the superclasses above; ArithValue is first so that
    e.g. __init__, and __str__ from ArithValue are used instead of
    from ScalarValue.
    """

    @cached_property
    def dtype(self) -> Type:
        return self.type

    @cached_property
    def literal_value(self) -> Union[int, float, bool]:
        if not self.is_constant():
            raise ValueError("Can't build literal from non-constant ScalarValue")
        return self.owner.opview.literal_value

    def __int__(self):
        return int(self.literal_value)

    def __float__(self):
        return float(self.literal_value)

    def coerce(self, other) -> Tuple["ScalarValue", "ScalarValue"]:
        if isinstance(other, (int, float, bool)):
            other = ScalarValue(other, dtype=self.dtype)
        elif isinstance(other, ScalarValue) and (
            isinstance(self.type, IndexType) or isinstance(other.type, IndexType)
        ):
            other = index_cast(other, to=self.type)
        else:
            raise ValueError(f"can't coerce {other=} to {self=}")
        return other


for t in [BF16Type, F16Type, F32Type, F64Type, IndexType, IntegerType, ComplexType]:
    register_value_caster(t.static_typeid)(ScalarValue)
