# 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.
# pylint: disable=no-else-return, unidiomatic-typecheck, invalid-name, arguments-differ
# ruff: noqa: RUF012
"""The expression functor of Relax."""

from collections.abc import Callable

import tvm_ffi

from tvm.ir import Call, Op
from tvm.ir.utils import derived_object
from tvm.runtime import Object

from ..ir.module import IRModule
from . import _ffi_api
from .block_builder import BlockBuilder
from .expr import (
    Binding,
    BindingBlock,
    Constant,
    DataflowBlock,
    DataflowVar,
    DataTypeImm,
    Expr,
    ExternFunc,
    Function,
    GlobalVar,
    If,
    MatchCast,
    SeqExpr,
    ShapeExpr,
    Span,
    StringImm,
    Tuple,
    TupleGetItem,
    Var,
    VarBinding,
)
from .type import Type

visitor = derived_object
"""
A decorator to wrap user-customized PyExprVisitor as TVM object _PyExprVisitor.

Parameters
----------
visitor_cls : PyExprVisitor
    The user-customized PyExprVisitor.

Returns
-------
cls : _PyExprVisitor
    The decorated TVM object _PyExprVisitor(ExprVisitor on the C++ side).

Example
-------
.. code-block:: python

    @relax.expr_functor.visitor
    class MyExprVisitor(PyExprVisitor):
        # customize visit function
        def visit_call_(self, op: Call) -> None:
            # just for demo purposes
            ...
    # myvisitor is now a special visitor that visit every Call with
    # user-customized visit_call_
    myvisitor = MyExprVisitor()
    # apply myvisitor to Expr/Binding/BindingBlock/VarDef
    myvisitor.visit_expr(expr)
    myvisitor.visit_binding(binding)
    myvisitor.visit_binding_block(bindingblock)
    myvisitor.visit_var_def(var)
"""

mutator = derived_object
"""
A decorator to wrap user-customized PyExprMutator as TVM object _PyExprMutator.
Note:  Cannot override visit function and post-order rewrite at the same time.

Parameters
----------
mutator_cls : PyExprMutator
    The user-customized PyExprMutator.

Returns
-------
cls : _PyExprMutator
    The decorated TVM object _PyExprMutator(ExprMutator on the C++ side).

Example
-------
.. code-block:: python

    @relax.expr_functor.mutator
    class MyExprMutator(PyExprMutator):
        # customize rewrite function
        def visit_tuple_(self, op: Tuple) -> Expr:
            # just for demo purposes
            ...

    # mymutator is now a special mutator that rewrite every Tuple with
    # user-customized visit_tuple_
    mymutator = MyExprMutator()
    # apply mymutator to Expr/Binding/BindingBlock/VarDef
    mymutator.visit_expr(expr)
    mymutator.visit_binding(binding)
    mymutator.visit_binding_block(bindingblock)
    mymutator.visit_var_def(var)
"""


class ExprFunctor:
    """
    An abstract visitor defined over Expr.
    Defines the default dispatch over expressions, and
    implements memoization.
    """

    def visit_expr(self, expr: Expr) -> Expr:
        """Apply the visitor to an expression."""
        if isinstance(expr, Constant):  # type: ignore
            ret = self.visit_constant_(expr)
        elif isinstance(expr, Tuple):
            ret = self.visit_tuple_(expr)
        elif isinstance(expr, DataflowVar):
            ret = self.visit_dataflow_var_(expr)
        elif isinstance(expr, Var):
            ret = self.visit_var_(expr)
        elif isinstance(expr, ShapeExpr):
            ret = self.visit_shape_expr_(expr)
        elif isinstance(expr, ExternFunc):
            ret = self.visit_extern_func_(expr)
        elif isinstance(expr, GlobalVar):  # type: ignore
            ret = self.visit_global_var_(expr)
        elif isinstance(expr, Function):
            ret = self.visit_function_(expr)
        elif isinstance(expr, Call):  # type: ignore
            ret = self.visit_call_(expr)
        elif isinstance(expr, SeqExpr):
            ret = self.visit_seq_expr_(expr)
        elif isinstance(expr, If):  # type: ignore
            ret = self.visit_if_(expr)
        elif isinstance(expr, Op):
            ret = self.visit_op_(expr)
        elif isinstance(expr, TupleGetItem):
            ret = self.visit_tuple_getitem_(expr)
        elif isinstance(expr, StringImm):
            ret = self.visit_string_imm_(expr)
        elif isinstance(expr, DataTypeImm):
            ret = self.visit_data_type_imm_(expr)
        elif isinstance(expr, Expr):
            ret = self.visit_expr_fallback_(expr)
        else:
            raise TypeError(f"Invalid type: {type(expr)}")

        return ret

    def visit_constant_(self, op: Constant):
        raise NotImplementedError()

    def visit_tuple_(self, op: Tuple):
        raise NotImplementedError()

    def visit_dataflow_var_(self, op: DataflowVar):
        raise NotImplementedError()

    def visit_var_(self, op: Var):
        raise NotImplementedError()

    def visit_shape_expr_(self, op: ShapeExpr):
        raise NotImplementedError()

    def visit_extern_func_(self, op: ExternFunc):
        raise NotImplementedError()

    def visit_global_var_(self, op: GlobalVar):
        raise NotImplementedError()

    def visit_function_(self, op: Function):
        raise NotImplementedError()

    def visit_call_(self, op: Call):
        raise NotImplementedError()

    def visit_seq_expr_(self, op: SeqExpr):
        raise NotImplementedError()

    def visit_if_(self, op: If):
        raise NotImplementedError()

    def visit_op_(self, op: Op):
        raise NotImplementedError()

    def visit_tuple_getitem_(self, op: TupleGetItem):
        raise NotImplementedError()

    def visit_expr_fallback_(self, op: Expr):
        raise NotImplementedError()

    def visit_string_imm_(self, op: StringImm):
        raise NotImplementedError()

    def visit_data_type_imm_(self, op: DataTypeImm):
        raise NotImplementedError()

    def visit_var_binding_(self, binding: VarBinding):
        raise NotImplementedError()

    def visit_match_cast_(self, binding: MatchCast):
        raise NotImplementedError()

    def visit_binding_block_(self, block: BindingBlock):
        raise NotImplementedError()

    def visit_dataflow_block_(self, block: DataflowBlock):
        raise NotImplementedError()

    def visit_var_def_(self, var: Var):
        raise NotImplementedError()

    def visit_dataflow_var_def_(self, var: DataflowVar):
        raise NotImplementedError()

    def visit_binding(self, binding: Binding):
        if isinstance(binding, MatchCast):
            self.visit_match_cast_(binding)
        elif isinstance(binding, VarBinding):
            self.visit_var_binding_(binding)
        else:
            raise TypeError(f"Invalid type: {type(binding)}")

    def visit_binding_block(self, block: BindingBlock):
        if isinstance(block, DataflowBlock):
            self.visit_dataflow_block_(block)
        elif isinstance(block, BindingBlock):
            self.visit_binding_block_(block)
        else:
            raise TypeError(f"Invalid type: {type(block)}")

    def visit_var_def(self, var: Var):
        if isinstance(var, DataflowVar):
            self.visit_dataflow_var_def_(var)
        elif isinstance(var, Var):
            self.visit_var_def_(var)
        else:
            raise TypeError(f"Invalid type: {type(var)}")


@tvm_ffi.register_object("expr_functor.PyExprVisitor")
class _PyExprVisitor(tvm_ffi.core.Object):
    """
    A TVM object to support customization of ExprVisitor on the python side.
    This is the decorated result returned from visitor decorator.

    WARNING: This is NOT the user facing class for method overwriting inheritance.

    See also: visitor, PyExprVisitor
    """

    def __init__(
        self,
        f_visit_expr: Callable | None = None,
        f_visit_constant_: Callable | None = None,
        f_visit_tuple_: Callable | None = None,
        f_visit_var_: Callable | None = None,
        f_visit_dataflow_var_: Callable | None = None,
        f_visit_shape_expr_: Callable | None = None,
        f_visit_extern_func_: Callable | None = None,
        f_visit_global_var_: Callable | None = None,
        f_visit_function_: Callable | None = None,
        f_visit_call_: Callable | None = None,
        f_visit_seq_expr_: Callable | None = None,
        f_visit_if_: Callable | None = None,
        f_visit_op_: Callable | None = None,
        f_visit_tuple_getitem_: Callable | None = None,
        f_visit_expr_fallback_: Callable | None = None,
        f_visit_string_imm_: Callable | None = None,
        f_visit_data_type_imm_: Callable | None = None,
        f_visit_binding: Callable | None = None,
        f_visit_var_binding_: Callable | None = None,
        f_visit_match_cast_: Callable | None = None,
        f_visit_binding_block: Callable | None = None,
        f_visit_binding_block_: Callable | None = None,
        f_visit_dataflow_block_: Callable | None = None,
        f_visit_var_def: Callable | None = None,
        f_visit_var_def_: Callable | None = None,
        f_visit_dataflow_var_def_: Callable | None = None,
        f_visit_span: Callable | None = None,
    ) -> None:
        """Constructor."""

        self.__init_handle_by_constructor__(
            _ffi_api.MakePyExprVisitor,  # type: ignore
            f_visit_expr,
            f_visit_constant_,
            f_visit_tuple_,
            f_visit_var_,
            f_visit_dataflow_var_,
            f_visit_shape_expr_,
            f_visit_extern_func_,
            f_visit_global_var_,
            f_visit_function_,
            f_visit_call_,
            f_visit_seq_expr_,
            f_visit_if_,
            f_visit_op_,
            f_visit_tuple_getitem_,
            f_visit_expr_fallback_,
            f_visit_string_imm_,
            f_visit_data_type_imm_,
            f_visit_binding,
            f_visit_var_binding_,
            f_visit_match_cast_,
            f_visit_binding_block,
            f_visit_binding_block_,
            f_visit_dataflow_block_,
            f_visit_var_def,
            f_visit_var_def_,
            f_visit_dataflow_var_def_,
            f_visit_span,
        )

    def visit_expr(self, expr: Expr) -> None:
        """Generic dispatcher for Expr.

        Parameters
        ----------
        expr : Expr
            The expr to be visited.
        """
        return _ffi_api.PyExprVisitorVisitExpr(self, expr)  # type: ignore

    def visit_binding(self, binding: Binding) -> None:
        """Generic dispatcher for Binding.

        Parameters
        ----------
        binding : Binding
            The binding to be visited.
        """
        return _ffi_api.PyExprVisitorVisitBinding(self, binding)  # type: ignore

    def visit_binding_block(self, block: BindingBlock) -> None:
        """Generic dispatcher for BindingBlock.

        Parameters
        ----------
        block : BindingBlock
            The block to be visited.
        """
        return _ffi_api.PyExprVisitorVisitBindingBlock(self, block)  # type: ignore

    def visit_var_def(self, var: Var) -> None:
        """Generic dispatcher for visiting the var definition site.
        Note that visit_var_() will only visit the usage site of an Var.

        Parameters
        ----------
        var : Var
            The var to be visited.
        """
        return _ffi_api.PyExprVisitorVisitVarDef(self, var)  # type: ignore


class PyExprVisitor:
    """
    An abstract ExprVisitor with customized methods on the python-side.
    This is the user facing class for method overwriting inheritance.
    _tvm_metadata discribes the class to inherit("cls"), the methods
    that users can overwrite("methods").

    Note: @relax.expr_functor.visitor is required for proper usage of any inherited class.

    See also: visitor, _PyExprVisitor

    Example:

    .. code-block:: python

        @relax.expr_functor.visitor
        def MyExprVisitor(PyExprVisitor):
            ...

    """

    _tvm_metadata = {
        "cls": _PyExprVisitor,
        "methods": [
            "visit_expr",
            "visit_constant_",
            "visit_tuple_",
            "visit_var_",
            "visit_dataflow_var_",
            "visit_shape_expr_",
            "visit_extern_func_",
            "visit_global_var_",
            "visit_function_",
            "visit_call_",
            "visit_seq_expr_",
            "visit_if_",
            "visit_op_",
            "visit_tuple_getitem_",
            "visit_expr_fallback_",
            "visit_string_imm_",
            "visit_data_type_imm_",
            "visit_binding",
            "visit_var_binding_",
            "visit_match_cast_",
            "visit_binding_block",
            "visit_binding_block_",
            "visit_dataflow_block_",
            "visit_var_def",
            "visit_var_def_",
            "visit_dataflow_var_def_",
            "visit_span",
        ],
    }

    def visit_expr(self, expr: Expr) -> None:
        """Generic dispatcher for Expr.
        Users can customized this function to overwrite VisitExpr(const Expr& expr) on the C++ side.

        Parameters
        ----------
        expr : Expr
            The expr to be visited.
        """
        # Using self._outer() to ref _PyExprVisitor
        return _ffi_api.PyExprVisitorVisitExpr(self._outer(), expr)  # type: ignore

    def visit_binding(self, binding: Binding) -> None:
        """Generic dispatcher for Binding.
        Users can customized this function to overwrite VisitBinding(const Binding& binding)
        on the C++ side.

        Parameters
        ----------
        binding : Binding
            The binding to be visited.
        """
        # Using self._outer() to ref _PyExprVisitor
        return _ffi_api.PyExprVisitorVisitBinding(self._outer(), binding)  # type: ignore

    def visit_binding_block(self, block: BindingBlock) -> None:
        """Generic dispatcher for BindingBlock.
        Users can customized this function to overwrite VisitBindingBlock(const BindingBlock& block)
        on the C++ side.

        Parameters
        ----------
        block : BindingBlock
            The block to be visited.
        """
        # Using self._outer() to ref _PyExprVisitor
        return _ffi_api.PyExprVisitorVisitBindingBlock(self._outer(), block)  # type: ignore

    def visit_var_def(self, var: Var) -> None:
        """Generic dispatcher for visiting the var definition site.
        Users can customized this function to overwrite VisitVarDef(const Var& var) on the C++ side.
        Note that visit_var_() will only visit the usage site of an Var.

        Parameters
        ----------
        var : Var
            The var to be visited.
        """
        # Using self._outer() to ref _PyExprVisitor
        return _ffi_api.PyExprVisitorVisitVarDef(self._outer(), var)  # type: ignore

    def visit_constant_(self, op: Constant) -> None:
        """Visit Constant.
        Users can customized this function to overwrite VisitExpr_(const ConstantNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Constant
            The Constant to be visited.
        """
        # Using self._outer() to ref _PyExprVisitor
        return _ffi_api.ExprVisitorVisitExpr(self._outer(), op)  # type: ignore

    def visit_tuple_(self, op: Tuple) -> None:
        """Visit Tuple.
        Users can customized this function to overwrite VisitExpr_(const TupleNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Tuple
            The Tuple to be visited.
        """
        # Using self._outer() to ref _PyExprVisitor
        return _ffi_api.ExprVisitorVisitExpr(self._outer(), op)  # type: ignore

    def visit_var_(self, op: Var) -> None:
        """Visit Var.
        Users can customized this function to overwrite VisitExpr_(const VarNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Var
            The Var to be visited.
        """
        # Using self._outer() to ref _PyExprVisitor
        return _ffi_api.ExprVisitorVisitExpr(self._outer(), op)  # type: ignore

    def visit_dataflow_var_(self, op: DataflowVar) -> None:
        """Visit DataflowVar.
        Users can customized this function to overwrite VisitExpr_(const DataflowVarNode* op)
        on the C++ side.

        Parameters
        ----------
        op : DataflowVar
            The DataflowVar to be visited.
        """
        # Using self._outer() to ref _PyExprVisitor
        return _ffi_api.ExprVisitorVisitExpr(self._outer(), op)  # type: ignore

    def visit_shape_expr_(self, op: ShapeExpr) -> None:
        """Visit ShapeExpr.
        Users can customized this function to overwrite VisitExpr_(const ShapeExprNode* op)
        on the C++ side.

        Parameters
        ----------
        op : ShapeExpr
            The ShapeExpr to be visited.
        """
        # Using self._outer() to ref _PyExprVisitor
        return _ffi_api.ExprVisitorVisitExpr(self._outer(), op)  # type: ignore

    def visit_extern_func_(self, op: ExternFunc) -> None:
        """Visit ExternFunc.
        Users can customized this function to overwrite VisitExpr_(const ExternFuncNode* op)
        on the C++ side.

        Parameters
        ----------
        op : ExternFunc
            The ExternFunc to be visited.
        """
        # Using self._outer() to ref _PyExprVisitor
        return _ffi_api.ExprVisitorVisitExpr(self._outer(), op)  # type: ignore

    def visit_global_var_(self, op: GlobalVar) -> None:
        """Visit GlobalVar.
        Users can customized this function to overwrite VisitExpr_(const GlobalVarNode* op)
        on the C++ side.

        Parameters
        ----------
        op : GlobalVar
            The GlobalVar to be visited.
        """
        # Using self._outer() to ref _PyExprVisitor
        return _ffi_api.ExprVisitorVisitExpr(self._outer(), op)  # type: ignore

    def visit_function_(self, op: Function) -> None:
        """Visit Function.
        Users can customized this function to overwrite VisitExpr_(const FunctionNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Function
            The Function to be visited.
        """
        # Using self._outer() to ref _PyExprVisitor
        return _ffi_api.ExprVisitorVisitExpr(self._outer(), op)  # type: ignore

    def visit_call_(self, op: Call) -> None:
        """Visit Call.
        Users can customized this function to overwrite VisitExpr_(const CallNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Call
            The Call to be visited.
        """
        # Using self._outer() to ref _PyExprVisitor
        return _ffi_api.ExprVisitorVisitExpr(self._outer(), op)  # type: ignore

    def visit_seq_expr_(self, op: SeqExpr) -> None:
        """Visit SeqExpr.
        Users can customized this function to overwrite VisitExpr_(const SeqExprNode* op)
        on the C++ side.

        Parameters
        ----------
        op : SeqExpr
            The SeqExpr to be visited.
        """
        # Using self._outer() to ref _PyExprVisitor
        return _ffi_api.ExprVisitorVisitExpr(self._outer(), op)  # type: ignore

    def visit_if_(self, op: If) -> None:
        """Visit If.
        Users can customized this function to overwrite VisitExpr_(const IfNode* op)
        on the C++ side.

        Parameters
        ----------
        op : If
            The If to be visited.
        """
        # Using self._outer() to ref _PyExprVisitor
        return _ffi_api.ExprVisitorVisitExpr(self._outer(), op)  # type: ignore

    def visit_op_(self, op: Op) -> None:
        """Visit Op.
        Users can customized this function to overwrite VisitExpr_(const OpNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Op
            The Op to be visited.
        """
        # Using self._outer() to ref _PyExprVisitor
        return _ffi_api.ExprVisitorVisitExpr(self._outer(), op)  # type: ignore

    def visit_tuple_getitem_(self, op: TupleGetItem) -> None:
        """Visit TupleGetItem.
        Users can customized this function to overwrite VisitExpr_(const TupleGetItemNode* op)
        on the C++ side.

        Parameters
        ----------
        op : TupleGetItem
            The TupleGetItem to be visited.
        """
        # Using self._outer() to ref _PyExprVisitor
        return _ffi_api.ExprVisitorVisitExpr(self._outer(), op)  # type: ignore

    def visit_expr_fallback_(self, op: Expr) -> None:
        """Visit an expression handled by the C++ fallback.

        Parameters
        ----------
        op : Expr
            The expression to be visited.
        """
        # Using self._outer() to ref _PyExprVisitor
        return _ffi_api.ExprVisitorVisitExprFallback(self._outer(), op)  # type: ignore

    def visit_string_imm_(self, op: StringImm) -> None:
        """Visit StringImm.
        Users can customized this function to overwrite VisitExpr_(const StringImmNode* op)
        on the C++ side.

        Parameters
        ----------
        op : StringImm
            The StringImm to be visited.
        """
        # Using self._outer() to ref _PyExprVisitor
        return _ffi_api.ExprVisitorVisitExpr(self._outer(), op)  # type: ignore

    def visit_data_type_imm_(self, op: DataTypeImm) -> None:
        """Visit DataTypeImm.
        Users can customized this function to overwrite VisitExpr_(const DataTypeImmNode* op)
        on the C++ side.

        Parameters
        ----------
        op : DataTypeImm
            The DataTypeImm to be visited.
        """
        # Using self._outer() to ref _PyExprVisitor
        return _ffi_api.ExprVisitorVisitExpr(self._outer(), op)  # type: ignore

    def visit_var_binding_(self, binding: VarBinding) -> None:
        """Visit VarBinding.
        Users can customized this function to overwrite VisitBinding_(const VarBindingNode* binding)
        on the C++ side.

        Parameters
        ----------
        binding : VarBinding
            The VarBinding to be visited.
        """
        # Using self._outer() to ref _PyExprVisitor
        return _ffi_api.ExprVisitorVisitBinding(self._outer(), binding)  # type: ignore

    def visit_match_cast_(self, binding: MatchCast) -> None:
        """Visit MatchCast.
        Users can customized this function to overwrite VisitBinding_(const MatchCastNode* binding)
        on the C++ side.

        Parameters
        ----------
        binding : MatchCast
            The MatchCast to be visited.
        """
        # Using self._outer() to ref _PyExprVisitor
        return _ffi_api.ExprVisitorVisitBinding(self._outer(), binding)  # type: ignore

    def visit_binding_block_(self, block: BindingBlock) -> None:
        """Visit BindingBlock.
        Users can customized this function to overwrite VisitBindingBlock_(const BindingBlockNode*
        block) on the C++ side.

        Parameters
        ----------
        block : BindingBlock
            The BindingBlock to be visited.
        """
        # Using self._outer() to ref _PyExprVisitor
        return _ffi_api.ExprVisitorVisitBindingBlock(self._outer(), block)  # type: ignore

    def visit_dataflow_block_(self, block: DataflowBlock) -> None:
        """Visit DataflowBlock.
        Users can customized this function to overwrite VisitBindingBlock_(const DataflowBlockNode*
        block) on the C++ side.

        Parameters
        ----------
        block : DataflowBlock
            The DataflowBlock to be visited.
        """
        # Using self._outer() to ref _PyExprVisitor
        return _ffi_api.ExprVisitorVisitBindingBlock(self._outer(), block)  # type: ignore

    def visit_var_def_(self, var: Var) -> None:
        """Visit the Var definition site.
        Users can customized this function to overwrite VisitVarDef_(const VarNode* var)
        on the C++ side.

        Parameters
        ----------
        var : Var
            The Var to be visited.
        """
        # Using self._outer() to ref _PyExprVisitor
        return _ffi_api.ExprVisitorVisitVarDef(self._outer(), var)  # type: ignore

    def visit_dataflow_var_def_(self, var: DataflowVar) -> None:
        """Visit the DataflowVar definition site.
        Users can customized this function to overwrite VisitVarDef_(const DataflowVarNode* var)
        on the C++ side.

        Parameters
        ----------
        var : DataflowVar
            The DataflowVar to be visited.
        """
        # Using self._outer() to ref _PyExprVisitor
        return _ffi_api.ExprVisitorVisitVarDef(self._outer(), var)  # type: ignore

    def visit_span(self, span: Span) -> None:
        """Visit Span.
        Users can customized this function to overwrite VisitSpan(const Span& span) on the C++ side.

        Parameters
        ----------
        span : Span
            The Span to be visited.
        """
        # Using self._outer() to ref _PyExprVisitor
        return _ffi_api.ExprVisitorVisitSpan(self._outer(), span)  # type: ignore


@tvm_ffi.register_object("expr_functor.PyExprMutator")
class _PyExprMutator(Object):
    """
    A TVM object to support customization of ExprMutator on the python side.
    This is the decorated result returned from mutator decorator.

    WARNING: This is NOT the user facing class for method overwriting inheritance.

    See also: mutator, PyExprmutator
    """

    def __init__(
        self,
        builder: BlockBuilder = None,
        f_visit_expr: Callable | None = None,
        f_visit_constant_: Callable | None = None,
        f_visit_tuple_: Callable | None = None,
        f_visit_var_: Callable | None = None,
        f_visit_dataflow_var_: Callable | None = None,
        f_visit_shape_expr_: Callable | None = None,
        f_visit_extern_func_: Callable | None = None,
        f_visit_global_var_: Callable | None = None,
        f_visit_function_: Callable | None = None,
        f_visit_call_: Callable | None = None,
        f_visit_seq_expr_: Callable | None = None,
        f_visit_if_: Callable | None = None,
        f_visit_op_: Callable | None = None,
        f_visit_tuple_getitem_: Callable | None = None,
        f_visit_expr_fallback_: Callable | None = None,
        f_visit_string_imm_: Callable | None = None,
        f_visit_data_type_imm_: Callable | None = None,
        f_visit_binding: Callable | None = None,
        f_visit_var_binding_: Callable | None = None,
        f_visit_match_cast_: Callable | None = None,
        f_visit_binding_block: Callable | None = None,
        f_visit_binding_block_: Callable | None = None,
        f_visit_dataflow_block_: Callable | None = None,
        f_visit_var_def: Callable | None = None,
        f_visit_var_def_: Callable | None = None,
        f_visit_dataflow_var_def_: Callable | None = None,
        f_visit_span: Callable | None = None,
    ) -> None:
        """Constructor."""

        self.__init_handle_by_constructor__(
            _ffi_api.MakePyExprMutator,  # type: ignore
            builder,
            f_visit_expr,
            f_visit_constant_,
            f_visit_tuple_,
            f_visit_var_,
            f_visit_dataflow_var_,
            f_visit_shape_expr_,
            f_visit_extern_func_,
            f_visit_global_var_,
            f_visit_function_,
            f_visit_call_,
            f_visit_seq_expr_,
            f_visit_if_,
            f_visit_op_,
            f_visit_tuple_getitem_,
            f_visit_expr_fallback_,
            f_visit_string_imm_,
            f_visit_data_type_imm_,
            f_visit_binding,
            f_visit_var_binding_,
            f_visit_match_cast_,
            f_visit_binding_block,
            f_visit_binding_block_,
            f_visit_dataflow_block_,
            f_visit_var_def,
            f_visit_var_def_,
            f_visit_dataflow_var_def_,
            f_visit_span,
        )

    def visit_expr(self, expr: Expr) -> Expr:
        """Generic dispatcher for Expr.

        Parameters
        ----------
        expr : Expr
            The expr to be visited.

        Returns
        -------
        result : Expr
            The Expr after transformation.
        """
        return _ffi_api.PyExprMutatorVisitExpr(self, expr)  # type: ignore

    def visit_binding(self, binding: Binding) -> None:
        """Generic dispatcher for Binding.

        Parameters
        ----------
        binding : Binding
            The binding to be visited.
        """
        return _ffi_api.PyExprMutatorVisitBinding(self, binding)  # type: ignore

    def visit_binding_block(self, block: BindingBlock) -> BindingBlock:
        """Generic dispatcher for BindingBlock.

        Parameters
        ----------
        block : BindingBlock
            The block to be visited.

        Returns
        -------
        result : BindingBlock
            The binding block after transformation.
        """
        return _ffi_api.PyExprMutatorVisitBindingBlock(self, block)  # type: ignore

    def visit_var_def(self, var: Var) -> Var:
        """Generic dispatcher for visiting the var definition site.
        Note that visit_var_() will only visit the usage site of an Var.

        Parameters
        ----------
        var : Var
            The var to be visited.

        Returns
        -------
        result : Var
            The var after post-order rewritten.
        """
        return _ffi_api.PyExprMutatorVisitVarDef(self, var)  # type: ignore


class PyExprMutator:
    """
    An abstract ExprMutator with customized methods on the python-side.
    This is the user facing class for method overwriting inheritance.
    _tvm_metadata discribes the class to inherit("cls"), the methods that users can
    overwrite("methods"), the constructor's parameters("fields")

    Note: @relax.expr_functor.mutator is required for proper usage of any inherited class.

    See also: visitor, _PyExprVisitor

    Example:

    .. code-block:: python

        @relax.expr_functor.mutator
        def MyExprMutator(PyExprMutator):
            ...

    """

    _tvm_metadata = {
        "cls": _PyExprMutator,
        "fields": ["builder_"],
        "methods": [
            "visit_expr",
            "visit_constant_",
            "visit_tuple_",
            "visit_var_",
            "visit_dataflow_var_",
            "visit_shape_expr_",
            "visit_extern_func_",
            "visit_global_var_",
            "visit_function_",
            "visit_call_",
            "visit_seq_expr_",
            "visit_if_",
            "visit_op_",
            "visit_tuple_getitem_",
            "visit_expr_fallback_",
            "visit_string_imm_",
            "visit_data_type_imm_",
            "visit_binding",
            "visit_var_binding_",
            "visit_match_cast_",
            "visit_binding_block",
            "visit_binding_block_",
            "visit_dataflow_block_",
            "visit_var_def",
            "visit_var_def_",
            "visit_dataflow_var_def_",
            "visit_span",
        ],
    }

    def __init__(self, mod: IRModule | None = None) -> None:
        """Constructor"""
        self.builder_ = BlockBuilder(mod)

    def visit_expr(self, expr: Expr) -> Expr:
        """Generic dispatcher for Expr.
        Users can customized this function to overwrite VisitExpr(const Expr& expr) on the C++ side.

        Parameters
        ----------
        expr : Expr
            The expr to be visited.

        Returns
        -------
        result : Expr
            The Expr after transformation.
        """
        # Using self._outer() to ref _PyExprMutator
        return _ffi_api.PyExprMutatorVisitExpr(self._outer(), expr)  # type: ignore

    def visit_binding(self, binding: Binding) -> None:
        """Generic dispatcher for Binding.
        Users can customized this function to overwrite VisitBinding(const Binding& binding)
        on the C++ side.

        Parameters
        ----------
        binding : Binding
            The binding to be visited.
        """
        # Using self._outer() to ref _PyExprMutator
        return _ffi_api.PyExprMutatorVisitBinding(self._outer(), binding)  # type: ignore

    def visit_binding_block(self, block: BindingBlock) -> BindingBlock:
        """Generic dispatcher for BindingBlock.
        Users can customized this function to overwrite VisitBindingBlock(const BindingBlock& block)
        on the C++ side.

        Parameters
        ----------
        block : BindingBlock
            The block to be visited.

        Returns
        -------
        result : BindingBlock
            The binding block after transformation.
        """
        # Using self._outer() to ref _PyExprMutator
        return _ffi_api.PyExprMutatorVisitBindingBlock(self._outer(), block)  # type: ignore

    def visit_var_def(self, var: Var) -> Var:
        """Generic dispatcher for visiting the var definition site.
        Users can customized this function to overwrite VisitVarDef(const Var& var) on the C++ side.
        Note that visit_var_() will only visit the usage site of an Var.

        Parameters
        ----------
        var : Var
            The var to be visited.

        Returns
        -------
        result: Var
            The var after post-order rewritten.
        """
        # Using self._outer() to ref _PyExprMutator
        return _ffi_api.PyExprMutatorVisitVarDef(self._outer(), var)  # type: ignore

    def visit_constant_(self, op: Constant) -> Expr:
        """Visit Constant.
        Users can customized this function to overwrite VisitExpr_(const ConstantNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Constant
            The Constant to be visited.

        Returns
        -------
        result : Expr
            The Expr after transformation
        """
        # Using self._outer() to ref _PyExprMutator
        return _ffi_api.ExprMutatorVisitExpr(self._outer(), op)  # type: ignore

    def visit_tuple_(self, op: Tuple) -> Expr:
        """Visit Tuple.
        Users can customized this function to overwrite VisitExpr_(const TupleNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Tuple
            The Tuple to be visited.

        Returns
        -------
        result : Expr
            The Expr after transformation
        """
        # Using self._outer() to ref _PyExprMutator
        return _ffi_api.ExprMutatorVisitExpr(self._outer(), op)  # type: ignore

    def visit_var_(self, op: Var) -> Expr:
        """Visit Var.
        Users can customized this function to overwrite VisitExpr_(const VarNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Var
            The Var to be visited.

        Returns
        -------
        result : Expr
            The Expr after transformation
        """
        # Using self._outer() to ref _PyExprMutator
        return _ffi_api.ExprMutatorVisitExpr(self._outer(), op)  # type: ignore

    def visit_dataflow_var_(self, op: DataflowVar) -> Expr:
        """Visit DataflowVar.
        Users can customized this function to overwrite VisitExpr_(const DataflowVarNode* op)
        on the C++ side.

        Parameters
        ----------
        op : DataflowVar
            The DataflowVar to be visited.

        Returns
        -------
        result : Expr
            The Expr after transformation
        """
        # Using self._outer() to ref _PyExprMutator
        return _ffi_api.ExprMutatorVisitExpr(self._outer(), op)  # type: ignore

    def visit_shape_expr_(self, op: ShapeExpr) -> Expr:
        """Visit ShapeExpr.
        Users can customized this function to overwrite VisitExpr_(const ShapeExprNode* op)
        on the C++ side.

        Parameters
        ----------
        op : ShapeExpr
            The ShapeExpr to be visited.

        Returns
        -------
        result : Expr
            The Expr after transformation
        """
        # Using self._outer() to ref _PyExprMutator
        return _ffi_api.ExprMutatorVisitExpr(self._outer(), op)  # type: ignore

    def visit_extern_func_(self, op: ExternFunc) -> Expr:
        """Visit ExternFunc.
        Users can customized this function to overwrite VisitExpr_(const ExternFuncNode* op)
        on the C++ side.

        Parameters
        ----------
        op : ExternFunc
            The ExternFunc to be visited.

        Returns
        -------
        result : Expr
            The Expr after transformation
        """
        # Using self._outer() to ref _PyExprMutator
        return _ffi_api.ExprMutatorVisitExpr(self._outer(), op)  # type: ignore

    def visit_global_var_(self, op: GlobalVar) -> Expr:
        """Visit GlobalVar.
        Users can customized this function to overwrite VisitExpr_(const GlobalVarNode* op)
        on the C++ side.

        Parameters
        ----------
        op : GlobalVar
            The GlobalVar to be visited.

        Returns
        -------
        result : Expr
            The Expr after transformation
        """
        # Using self._outer() to ref _PyExprMutator
        return _ffi_api.ExprMutatorVisitExpr(self._outer(), op)  # type: ignore

    def visit_function_(self, op: Function) -> Expr:
        """Visit Function.
        Users can customized this function to overwrite VisitExpr_(const FunctionNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Function
            The Function to be visited.

        Returns
        -------
        result : Expr
            The Expr after transformation
        """
        # Using self._outer() to ref _PyExprMutator
        return _ffi_api.ExprMutatorVisitExpr(self._outer(), op)  # type: ignore

    def visit_call_(self, op: Call) -> Expr:
        """Visit Call.
        Users can customized this function to overwrite VisitExpr_(const CallNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Call
            The Call to be visited.

        Returns
        -------
        result : Expr
            The Expr after transformation
        """
        # Using self._outer() to ref _PyExprMutator
        return _ffi_api.ExprMutatorVisitExpr(self._outer(), op)  # type: ignore

    def visit_seq_expr_(self, op: SeqExpr) -> Expr:
        """Visit SeqExpr.
        Users can customized this function to overwrite VisitExpr_(const SeqExprNode* op)
        on the C++ side.

        Parameters
        ----------
        op : SeqExpr
            The SeqExpr to be visited.

        Returns
        -------
        result : Expr
            The Expr after transformation
        """
        # Using self._outer() to ref _PyExprMutator
        return _ffi_api.ExprMutatorVisitExpr(self._outer(), op)  # type: ignore

    def visit_if_(self, op: If) -> Expr:
        """Visit If.
        Users can customized this function to overwrite VisitExpr_(const IfNode* op)
        on the C++ side.

        Parameters
        ----------
        op : If
            The If to be visited.

        Returns
        -------
        result : Expr
            The Expr after transformation
        """
        # Using self._outer() to ref _PyExprMutator
        return _ffi_api.ExprMutatorVisitExpr(self._outer(), op)  # type: ignore

    def visit_op_(self, op: Op) -> Expr:
        """Visit Op.
        Users can customized this function to overwrite VisitExpr_(const OpNode* op)
        on the C++ side.

        Parameters
        ----------
        op : Op
            The Op to be visited.

        Returns
        -------
        result : Expr
            The Expr after transformation
        """
        # Using self._outer() to ref _PyExprMutator
        return _ffi_api.ExprMutatorVisitExpr(self._outer(), op)  # type: ignore

    def visit_tuple_getitem_(self, op: TupleGetItem) -> Expr:
        """Visit TupleGetItem.
        Users can customized this function to overwrite VisitExpr_(const TupleGetItemNode* op)
        on the C++ side.

        Parameters
        ----------
        op : TupleGetItem
            The TupleGetItem to be visited.

        Returns
        -------
        result : Expr
            The Expr after transformation
        """
        # Using self._outer() to ref _PyExprMutator
        return _ffi_api.ExprMutatorVisitExpr(self._outer(), op)  # type: ignore

    def visit_expr_fallback_(self, op: Expr) -> Expr:
        """Visit an expression handled by the C++ fallback.

        Parameters
        ----------
        op : Expr
            The expression to be visited.

        Returns
        -------
        result : Expr
            The Expr after transformation
        """
        # Using self._outer() to ref _PyExprMutator
        return _ffi_api.ExprMutatorVisitExprFallback(self._outer(), op)  # type: ignore

    def visit_string_imm_(self, op: StringImm) -> Expr:
        """Visit StringImm.
        Users can customized this function to overwrite VisitExpr_(const StringImmNode* op)
        on the C++ side.

        Parameters
        ----------
        op : StringImm
            The StringImm to be visited.

        Returns
        -------
        result : Expr
            The Expr after transformation
        """
        # Using self._outer() to ref _PyExprMutator
        return _ffi_api.ExprMutatorVisitExpr(self._outer(), op)  # type: ignore

    def visit_data_type_imm_(self, op: DataTypeImm) -> Expr:
        """Visit DataTypeImm.
        Users can customized this function to overwrite VisitExpr_(const DataTypeImmNode* op)
        on the C++ side.

        Parameters
        ----------
        op : DataTypeImm
            The DataTypeImm to be visited.

        Returns
        -------
        result : Expr
            The Expr after transformation
        """
        # Using self._outer() to ref _PyExprMutator
        return _ffi_api.ExprMutatorVisitExpr(self._outer(), op)  # type: ignore

    def visit_var_binding_(self, binding: VarBinding) -> None:
        """Visit VarBinding.
        Users can customized this function to overwrite VisitBinding_(const VarBindingNode* binding)
        on the C++ side.

        Parameters
        ----------
        binding : VarBinding
            The VarBinding to be visited.
        """
        # Using self._outer() to ref _PyExprMutator
        return _ffi_api.ExprMutatorVisitBinding(self._outer(), binding)  # type: ignore

    def visit_match_cast_(self, binding: MatchCast) -> None:
        """Visit MatchCast.
        Users can customized this function to overwrite VisitBinding_(const MatchCastNode* binding)
        on the C++ side.

        Parameters
        ----------
        binding : MatchCast
            The MatchCast to be visited.
        """
        # Using self._outer() to ref _PyExprMutator
        return _ffi_api.ExprMutatorVisitBinding(self._outer(), binding)  # type: ignore

    def visit_binding_block_(self, block: BindingBlock) -> BindingBlock:
        """Visit BindingBlock.
        Users can customized this function to overwrite VisitBindingBlock_(const BindingBlockNode*
        block) on the C++ side.

        Parameters
        ----------
        block : BindingBlock
            The BindingBlock to be visited.

        Returns
        -------
        result : BindingBlock
            The binding block after transformation
        """
        # Using self._outer() to ref _PyExprMutator
        return _ffi_api.ExprMutatorVisitBindingBlock(self._outer(), block)  # type: ignore

    def visit_dataflow_block_(self, block: DataflowBlock) -> BindingBlock:
        """Visit DataflowBlock.
        Users can customized this function to overwrite VisitBindingBlock_(const DataflowBlockNode*
        block) on the C++ side.

        Parameters
        ----------
        block : DataflowBlock
            The DataflowBlock to be visited.

        Returns
        -------
        result : BindingBlock
            The binding block after transformation
        """
        # Using self._outer() to ref _PyExprMutator
        return _ffi_api.ExprMutatorVisitBindingBlock(self._outer(), block)  # type: ignore

    def visit_var_def_(self, var: Var) -> Var:
        """Visit the Var definition site.
        Users can customized this function to overwrite VisitVarDef_(const VarNode* var)
        on the C++ side.

        Parameters
        ----------
        var : Var
            The Var to be visited.

        Returns
        -------
        result : Var
            The var after post-order rewritten.
        """
        # Using self._outer() to ref _PyExprMutator
        return _ffi_api.ExprMutatorVisitVarDef(self._outer(), var)  # type: ignore

    def visit_dataflow_var_def_(self, var: DataflowVar) -> Var:
        """Visit the DataflowVar definition site.
        Users can customized this function to overwrite VisitVarDef_(const DataflowVarNode* var)
        on the C++ side.

        Parameters
        ----------
        var : DataflowVar
            The DataflowVar to be visited.

        Returns
        -------
        result : Var
            The var after post-order rewritten.
        """
        # Using self._outer() to ref _PyExprMutator
        return _ffi_api.ExprMutatorVisitVarDef(self._outer(), var)  # type: ignore

    def visit_span(self, span: Span) -> Span:
        """Visit Span.
        Users can customized this function to overwrite VisitSpan(const Span& span) on the C++ side.

        Parameters
        ----------
        span : Span
            The Span to be visited.

        Returns
        -------
        result : Span
            The span after transformation.
        """
        raise NotImplementedError

    def visit_expr_post_order(self, expr: Expr) -> Expr:
        """Post-order rewrite an Expr and normalize.

        Parameters
        ----------
        expr : Expr
            The Expr to be rewritten.

        Returns
        -------
        result : Expr
            The Expr after post-order rewritten.
        """
        return _ffi_api.PyExprMutatorVisitExprPostOrder(self._outer(), expr)  # type: ignore

    def set_var_remap(self, old_var: Var, new_var: Var) -> None:
        """Remap a var to a new var in use-site.

        Parameters
        ----------
        old_var : Var
            The old var.
        new_var : Var
            The new var.
        """
        # Using self._outer() to ref _PyExprMutator
        return _ffi_api.PyExprMutatorSetVarRemap(  # type: ignore
            self._outer(), old_var, new_var
        )

    def get_var_remap(self, var: Var) -> Var:
        """Remap a var to a new var in use-site.

        Parameters
        ----------
        var : Var
            The old var.

        Returns
        -------
        var : Var
            The remapped var.
        """
        # Using self._outer() to ref _PyExprMutator
        return _ffi_api.PyExprMutatorGetVarRemap(self._outer(), var)  # type: ignore

    def visit_with_new_scope(self, expr: Expr) -> Expr:
        """Rewrite the expr with a new scope, used in a Function's body and the branches of If.

        Parameters
        ----------
        expr : Expr
            The expr to be visited.

        Returns
        -------
        var : Var
            The expr after visiting.
        """
        # Using self._outer() to ref _PyExprMutator
        return _ffi_api.PyExprMutatorVisitWithNewScope(self._outer(), expr)  # type: ignore

    def lookup_binding(self, var: Var) -> Expr | None:
        """Look up the value bound to a variable.
        Note: For function parameters, this function returns std::nullopt.

        Parameters
        ----------
        var : Var
            The var to be looked up.

        Returns
        -------
        var : Var
            The value bound to the input var.
        """
        # Using self._outer() to ref _PyExprMutator
        return _ffi_api.PyExprMutatorLookupBinding(self._outer(), var)  # type: ignore

    def with_type(self, var: Var, ty: Type) -> Var:
        """Create a new var with specified shape and type if the original var's shape or type does
        not match with the specified ones.

        Parameters
        ----------
        var : Var
            The var to be updated.
        ty : Type
            The type.

        Returns
        -------
        var : Var
            The var filled with shape and type.
        """
        # Using self._outer() to ref _PyExprMutator
        return _ffi_api.PyExprMutatorWithType(self._outer(), var, ty)  # type: ignore
