# Copyright (c) 2021 Graphcore Ltd. All rights reserved.
from typing import List, Union
import popart._internal.ir as _ir
from popxl.context import get_current_context, op_debug_context
from popxl.tensor import Tensor
from .utils import check_in_graph, handle_negative_axis


@op_debug_context
def split(t: Tensor, splits: Union[int, List[int]], axis: int = 0) -> List[Tensor]:
    """Split a tensor along an axis into a list of tensors.

    See also `PyTorch Tensor.split <https://pytorch.org/docs/stable/generated/torch.Tensor.split.html>`__, `NumPy split <https://numpy.org/doc/stable/reference/generated/numpy.split.html>`__, `ONNX Split <https://github.com/onnx/onnx/blob/main/docs/Operators.md#Split>`__.

    Args:
        t (Tensor): Tensor to be split.
        splits (Union[int, List[int]]): Either an int which specifies the
            number of splits or a list of ints specifying the length of
            each output tensor.
        axis (int, optional): Axis to split along. Defaults to 0.

    Raises:
        ValueError: If the split doesn't equally divide the tensor.

    Returns:
        List[Tensor]: A list of tensors.
    """
    ctx = get_current_context()
    g = ctx.graph
    pb_g = g._pb_graph

    check_in_graph(g, t=t)

    axis = handle_negative_axis(t, axis)

    if isinstance(splits, int):
        axis_len = t.shape[axis]
        if axis_len % splits != 0:
            raise ValueError(
                f"Split {splits} does not equally divide tensor axis {axis} of length {axis_len}."
            )
        splits = [axis_len // splits] * splits

    outputs_t = {
        i: g._create_tensor_id(f"{t.name}_split_{i}") for i in range(len(splits))
    }

    settings = ctx._get_op_settings("split")
    opid = _ir.OperatorIdentifier("ai.onnx", "Split", 2, _ir.NumInputs(1, 1), 1)
    op = pb_g.createConnectedOp_SplitOp(
        {0: t.id},
        outputs_t,
        axis_=axis,
        split_=splits,
        opid=opid,
        settings=settings,
    )

    output = [Tensor._from_pb_tensor(op.outTensor(i)) for i in range(len(splits))]

    return output