# Copyright (c) OpenMMLab. All rights reserved.
from typing import Union

import torch
from packaging import version

from juxtapose.mmdeploy.core import FUNCTION_REWRITER
from juxtapose.mmdeploy.utils.constants import Backend


# TODO add version control when MOD is supported by TensorRT
@FUNCTION_REWRITER.register_rewriter(
    func_name="torch.Tensor.__mod__", backend=Backend.TENSORRT.value
)
def mod__tensorrt(
    input: torch.Tensor, other: Union[torch.Tensor, torch.NumberType], *args, **kwargs
) -> torch.Tensor:
    """Rewrite `mod` when exporting model to ONNX for TensorRT backend."""
    ctx = FUNCTION_REWRITER.get_context()
    if version.parse(torch.__version__) > version.parse("1.10.0"):
        return input - (input // other) * other
    return ctx.origin_func(input, other, *args, **kwargs)
