# Copyright (c) OpenMMLab. All rights reserved.
import torch

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


@FUNCTION_REWRITER.register_rewriter(
    func_name="torch.Tensor.clip", backend=Backend.COREML.value
)
@FUNCTION_REWRITER.register_rewriter(
    func_name="torch.clip", backend=Backend.COREML.value
)
@FUNCTION_REWRITER.register_rewriter(
    func_name="torch.Tensor.clamp", backend=Backend.COREML.value
)
@FUNCTION_REWRITER.register_rewriter(
    func_name="torch.clamp", backend=Backend.COREML.value
)
def clip__coreml(input, min=None, max=None, **kwargs) -> torch.Tensor:
    """Rewrite `clip` for coreml backend.

    Cast data type.
    """
    ctx = FUNCTION_REWRITER.get_context()
    if min is not None and not isinstance(min, torch.Tensor):
        min = input.new_tensor(min)

    if max is not None and not isinstance(max, torch.Tensor):
        max = input.new_tensor(max)

    return ctx.origin_func(input, min=min, max=max, **kwargs)
