# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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
#
#   http://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.

"""Implementation of TIR operator."""

from tvm.ir import Op
from tvm.tirx import Expr
from tvm.tirx.stmt import TilePrimitiveCall


def get_tirx_op(op_name: str):
    assert isinstance(op_name, str)
    return Op.get("tirx.tile." + op_name)


class ArgProperty:
    def __init__(self, index):
        self.index = index

    def __get__(self, obj, objtype=None):
        assert obj is not None, "TilePrimitiveCall cannot be None"
        return obj.args[self.index]


### Base Operator Classes ###
class UnaryOp(TilePrimitiveCall):
    """Base class for unary operators: unary(output, input).

    Unary operators take a single input tensor and produce a single output tensor.
    """

    scalar_input = False
    output = ArgProperty(0)
    input = ArgProperty(1)

    @property
    def srcs(self) -> list[Expr]:
        """Get the source expression (input) of the operator."""
        return [self.input]

    @property
    def dsts(self) -> list[Expr]:
        """Get the destination expression (output) of the operator."""
        return [self.output]


class UnaryOpWithBiasScale(UnaryOp):
    """Extended unary operator with bias and scale parameters: unary_with_bias_scale(output, input, bias, scale).

    These operators support additional bias and scale parameters for more complex operations (only on trn).
    output = unary(input * scale + bias)
    """  # noqa: E501

    bias = ArgProperty(2)
    scale = ArgProperty(3)

    @property
    def srcs(self) -> list[Expr]:
        """Get the source expressions (inputs) of the operator."""
        return [self.input, self.bias, self.scale]


class BinaryOp(TilePrimitiveCall):
    """Base class for binary operators: binary(output, input0, input1).

    Binary operators take two input tensors and produce a single output tensor.
    """

    lhs = ArgProperty(1)
    rhs = ArgProperty(2)
    output = ArgProperty(0)

    @property
    def srcs(self) -> list[Expr]:
        """Get the source expressions (inputs) of the operator."""
        return [self.lhs, self.rhs]

    @property
    def dsts(self) -> list[Expr]:
        """Get the destination expression (output) of the operator."""
        return [self.output]


class ReduceOp(TilePrimitiveCall):
    """Base class for reduction operators: reduce(output, input, reduce_axes, accum).

    Reduction operators reduce one or more dimensions of the input tensor.
    """

    input = ArgProperty(1)
    output = ArgProperty(0)
    reduce_axes = ArgProperty(2)
    accum = ArgProperty(3)

    @property
    def srcs(self) -> list[Expr]:
        """Get the source expression (input) of the operator."""
        return [self.input]

    @property
    def dsts(self) -> list[Expr]:
        """Get the destination expression (output) of the operator."""
        return [self.output]


### Schedule Operators ###
class Zero(UnaryOp):
    """Zero out all elements in src and store to dst."""

    op = get_tirx_op("zero")


class Sqrt(UnaryOpWithBiasScale):
    """Compute square root of all elements in src and store to dst.

    If bias and scale are provided: dst = sqrt(src * scale + bias)
    """

    op = get_tirx_op("sqrt")


class Fill(UnaryOp):
    """Fill dst with a scalar value."""

    op = get_tirx_op("fill")
    scalar_input = True


class Add(BinaryOp):
    """Add src1 and src2 element-wise and store to dst."""

    op = get_tirx_op("add")


class Sub(BinaryOp):
    """Subtract src2 from src1 element-wise and store to dst."""

    op = get_tirx_op("sub")


class Mul(BinaryOp):
    """Multiply src1 and src2 element-wise and store to dst."""

    op = get_tirx_op("mul")


class FDiv(BinaryOp):
    """Divide src1 by src2 element-wise using floating point division and store to dst."""

    op = get_tirx_op("fdiv")


class FMA(TilePrimitiveCall):
    """Fused multiply-add: output = input * scale + bias.

    fma(output, input, scale, bias)

    scale and bias can each be either a BufferRegion or a Expr scalar.
    """

    op = get_tirx_op("fma")

    output = ArgProperty(0)
    input = ArgProperty(1)
    scale = ArgProperty(2)
    bias = ArgProperty(3)

    @property
    def srcs(self) -> list[Expr]:
        """Get the source expressions (inputs) of the operator."""
        return [self.input, self.scale, self.bias]

    @property
    def dsts(self) -> list[Expr]:
        """Get the destination expression (output) of the operator."""
        return [self.output]


class Cast(UnaryOp):
    """Cast src to dst."""

    op = get_tirx_op("cast")


class Copy(TilePrimitiveCall):
    """Copy all elements from src to dst.

    Args:
        dst: Destination buffer region
        src: Source buffer region
    """

    op = get_tirx_op("copy")

    dst = ArgProperty(0)
    src = ArgProperty(1)

    @property
    def srcs(self) -> list[Expr]:
        """Get the source expressions (inputs) of the operator."""
        return [self.src]

    @property
    def dsts(self) -> list[Expr]:
        """Get the destination expressions (outputs) of the operator."""
        return [self.dst]


class CopyAsync(TilePrimitiveCall):
    """Copy all elements from src to dst asynchronously.

    Args:
        dst: Destination buffer region
        src: Source buffer region
    """

    op = get_tirx_op("copy_async")

    dst = ArgProperty(0)
    src = ArgProperty(1)

    @property
    def srcs(self) -> list[Expr]:
        """Get the source expressions (inputs) of the operator."""
        return [self.src]

    @property
    def dsts(self) -> list[Expr]:
        """Get the destination expressions (outputs) of the operator."""
        return [self.dst]


class Gemm(TilePrimitiveCall):
    """General matrix multiplication: D = A * B * alpha + C * beta.

    Args:
        D: Output matrix
        A: First input matrix
        B: Second input matrix
        C: Third input matrix (for bias)
        transpose_A: Whether to transpose A
        transpose_B: Whether to transpose B
        alpha: Scalar multiplier for A*B
        beta: Scalar multiplier for C
    """

    op = get_tirx_op("gemm")
    output = ArgProperty(0)
    lhs = ArgProperty(1)
    rhs = ArgProperty(2)
    bias = ArgProperty(3)
    transpose_A = ArgProperty(4)
    transpose_B = ArgProperty(5)
    alpha = ArgProperty(6)
    beta = ArgProperty(7)

    @property
    def srcs(self) -> list[Expr]:
        """Get the source matrices."""
        return [self.lhs, self.rhs, self.bias]

    @property
    def dsts(self) -> list[Expr]:
        """Get the destination matrix."""
        return [self.output]


class GemmAsync(TilePrimitiveCall):
    """General matrix multiplication asynchronously.

    Supports two arg layouts:
    - Regular (6 args): C, A, B, transA, transB, accum
    - Block-scaled (8 args): C, A, B, SFA, SFB, transA, transB, accum
    """

    op = get_tirx_op("gemm_async")
    output = ArgProperty(0)
    lhs = ArgProperty(1)
    rhs = ArgProperty(2)

    @property
    def is_block_scaled(self) -> bool:
        """Whether this is a block-scaled MMA operation."""
        return len(self.args) == 8

    @property
    def sfa(self):
        """Get the scale factor buffer for A (None for regular MMA)."""
        return self.args[3] if self.is_block_scaled else None

    @property
    def sfb(self):
        """Get the scale factor buffer for B (None for regular MMA)."""
        return self.args[4] if self.is_block_scaled else None

    @property
    def transA(self):
        return self.args[5] if self.is_block_scaled else self.args[3]

    @property
    def transB(self):
        return self.args[6] if self.is_block_scaled else self.args[4]

    @property
    def accum(self):
        return self.args[7] if self.is_block_scaled else self.args[5]

    @property
    def srcs(self) -> list[Expr]:
        """Get the source matrices (including scale factors if block-scaled)."""
        srcs = [self.lhs, self.rhs]
        if self.is_block_scaled:
            srcs.extend([self.sfa, self.sfb])
        return srcs

    @property
    def dsts(self) -> list[Expr]:
        """Get the destination matrix."""
        return [self.output]


class Sum(ReduceOp):
    """Sum elements in src along specified axes and store in dst."""

    op = get_tirx_op("sum")


class Max(ReduceOp):
    """Compute maximum value in src along specified axes and store in dst."""

    op = get_tirx_op("max")


class Min(ReduceOp):
    """Compute minimum value in src along specified axes and store in dst."""

    op = get_tirx_op("min")


class Reciprocal(UnaryOp):
    """Compute reciprocal (1/x) for all elements in src and store to dst."""

    op = get_tirx_op("reciprocal")


class SiLU(UnaryOp):
    """Compute SiLU (x * sigmoid(x)) for all elements in src and store to dst."""

    op = get_tirx_op("silu")


class Memset(UnaryOp):
    """Set all elements in dst to a specified value."""

    op = get_tirx_op("memset")
    scalar_input = True


class Maximum(BinaryOp):
    """Compute element-wise maximum of src1 and src2 and store to dst."""

    op = get_tirx_op("maximum")


class Minimum(BinaryOp):
    """Compute element-wise minimum of src1 and src2 and store to dst."""

    op = get_tirx_op("minimum")


class Exp(UnaryOpWithBiasScale):
    """Compute exponential (e^x) of all elements in src and store to dst.

    If bias and scale are provided: dst = exp(src * scale + bias)
    """

    op = get_tirx_op("exp")


class Exp2(UnaryOpWithBiasScale):
    """Compute base-2 exponential (2^x) of all elements in src and store to dst.

    If bias and scale are provided: dst = exp2(src * scale + bias)
    """

    op = get_tirx_op("exp2")


class Select(BinaryOp):
    """Select elements from src1 or src2 based on the predicate.

    select(dst, src1, src2, predicate)
    """

    op = get_tirx_op("select")
    predicate = ArgProperty(3)


### Compose Ops ###
class BinaryReduce(TilePrimitiveCall):
    """Combine a binary operation with a reduction operation.

    binary_reduce(binary_output, reduce_output, binary_input1, binary_input2, binary_op, reduce_op, reduce_axes, )
    """  # noqa: E501

    op = get_tirx_op("binary_reduce")

    binary_output = ArgProperty(0)
    reduce_output = ArgProperty(1)
    binary_input1 = ArgProperty(2)
    binary_input2 = ArgProperty(3)
    binary_op = ArgProperty(4)
    reduce_op = ArgProperty(5)
    reduce_axes = ArgProperty(6)

    @property
    def srcs(self) -> list[Expr]:
        """Get the source expressions (inputs) of the operator."""
        return [self.binary_input1, self.binary_input2]

    @property
    def dsts(self) -> list[Expr]:
        """Get the destination expressions (outputs) of the operator."""
        return [self.binary_output, self.reduce_output]


class UnaryReduce(TilePrimitiveCall):
    """Combine a unary operation with a reduction operation.

    unary_reduce(unary_output, reduce_output, unary_input, unary_op, reduce_op, bias, scale, reduce_axes)
    """  # noqa: E501

    op = get_tirx_op("unary_reduce")

    unary_output = ArgProperty(0)
    reduce_output = ArgProperty(1)
    unary_input = ArgProperty(2)
    unary_op = ArgProperty(3)
    reduce_op = ArgProperty(4)
    bias = ArgProperty(5)
    scale = ArgProperty(6)
    reduce_axes = ArgProperty(7)

    @property
    def srcs(self) -> list[Expr]:
        """Get the source expressions (inputs) of the operator."""
        return [self.unary_input, self.bias, self.scale]

    @property
    def dsts(self) -> list[Expr]:
        """Get the destination expressions (outputs) of the operator."""
        return [self.unary_output, self.reduce_output]


class BinaryChain(TilePrimitiveCall):
    """Chain multiple binary operations together.

    binary_chain(output, data, operand0, operand1, op0, op1, reverse1)

    if not reverse1:
        output = (operand0 op0 data) op1 operand1
    else:
        output = operand1 op1 (operand0 op0 data)
    """

    op = get_tirx_op("binary_chain")

    output = ArgProperty(0)
    data = ArgProperty(1)
    operand0 = ArgProperty(2)
    operand1 = ArgProperty(3)
    op0 = ArgProperty(4)
    op1 = ArgProperty(5)
    reverse1 = ArgProperty(6)

    @property
    def srcs(self) -> list[Expr]:
        """Get the source expressions (inputs) of the operator."""
        return [self.data, self.operand0, self.operand1]

    @property
    def dsts(self) -> list[Expr]:
        """Get the destination expressions (outputs) of the operator."""
        return [self.output]


class ReduceNegate(ReduceOp):
    """
    Negate the result of a reduction operation.

    reduce_negate(output, input, reduce_axes, accum, reduce_op)
    """

    op = get_tirx_op("reduce_negate")

    reduce_op = ArgProperty(4)


class ComposeOp(TilePrimitiveCall):
    """Generic operator for composition of multiple operations.

    Must be lowered to specific compose operations before operator-level passes.
    """

    # TODO: add a pass to lower generic compose_op to specific compose ops

    op = get_tirx_op("compose_op")

    @property
    def srcs(self) -> list[Expr]:
        """Get the source expressions (inputs) of the operator."""
        raise NotImplementedError(
            "Generic compose_op must be lowered to specific compose ops before operator-level passes"  # noqa: E501
        )

    @property
    def dsts(self) -> list[Expr]:
        """Get the destination expressions (outputs) of the operator."""
        raise NotImplementedError(
            "Generic compose_op must be lowered to specific compose ops before operator-level passes"  # noqa: E501
        )


class PermuteLayout(TilePrimitiveCall):
    """Move data so the buffer's bytes are arranged under a different layout.

    Logical shape is preserved; only the byte placement changes. ``dst`` and
    ``src`` carry their own TileLayouts; on lowering, the dispatcher reads
    those layouts and emits a register-staged warp transpose, optionally
    inserting a bank-conflict-avoiding XOR-swizzle on the per-lane register
    slots.

    Args: ``permute_layout(dst_region, src_region)``.
    ``dst`` and ``src`` may alias the same underlying SMEM (in-place).
    """

    op = get_tirx_op("permute_layout")

    @property
    def dst(self) -> Expr:
        return self.args[0]

    @property
    def src(self) -> Expr:
        return self.args[1]

    @property
    def srcs(self) -> list[Expr]:
        return [self.src]

    @property
    def dsts(self) -> list[Expr]:
        return [self.dst]