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

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


@FUNCTION_REWRITER.register_rewriter(func_name="torch.Tensor.chunk", backend="ncnn")
def chunk__ncnn(self, num_chunks: int, dim: int = 0) -> torch.Tensor:
    """Rewrite `chunk` for NCNN backend.

    Chunk in ncnn are not supported, so it should be rewritten.
    """
    dim_len = self.shape[dim]
    # int ceil.
    step = dim_len // num_chunks
    if dim_len % num_chunks > 0:
        step += 1
    index_list = []
    index = 0
    while index < dim_len:
        index_list.append(index)
        index += step
    index_list.append(dim_len)
    output = [
        self.index_select(
            dim,
            torch.tensor(
                [j for j in range(index_list[i], index_list[i + 1])], dtype=torch.int64
            ),
        )
        for i in range(len(index_list) - 1)
    ]

    return output


@FUNCTION_REWRITER.register_rewriter(func_name="torch.Tensor.chunk", ir=IR.TORCHSCRIPT)
def chunk__torchscript(self, num_chunks: int, dim: int = 0) -> torch.Tensor:
    """Rewrite `chunk` for Torchscript.

    Replace chunk op with split op
    """
    dim_size = self.shape[dim]
    assert dim_size % num_chunks == 0, "cannot split to equal sizes"
    output = self.split(dim_size // num_chunks, dim=dim)
    return output
