# 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.flatten", backend=Backend.NCNN.value
)
@FUNCTION_REWRITER.register_rewriter(
    func_name="torch.flatten", backend=Backend.NCNN.value
)
@FUNCTION_REWRITER.register_rewriter(
    func_name="torch.Tensor.flatten", backend=Backend.COREML.value
)
@FUNCTION_REWRITER.register_rewriter(
    func_name="torch.flatten", backend=Backend.COREML.value
)
def flatten__coreml(input, start_dim=0, end_dim=-1) -> torch.Tensor:
    """Rewrite `flatten` for coreml backend.

    Use reshape instead of flatten
    """
    shape = input.shape
    end_dim = end_dim if end_dim > 0 else len(shape) + end_dim
    shape1 = list(shape[:start_dim])
    shape3 = list(shape[end_dim + 1 :])
    return input.reshape(shape1 + [-1] + shape3)
